TransformerLM_mathematical_definition_with_code_and_examples.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "TransformerLM - mathematical definition with code and examples.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "UitH_YdEEH5L"
      ],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/MichalRyszardWojcik/transformer-language-model/blob/main/TransformerLM_mathematical_definition_with_code_and_examples.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The transformer language model technology is used in the following projects:\n",
        "<ul>\n",
        "<li><a href=\"https://en.wikipedia.org/wiki/GPT-3\">GPT-3</a>: an AI system than can create articles, poetry, stories, news reports and dialogue using just a small amount of input text\n",
        "<li><a href='https://openai.com/blog/dall-e/'>DALL·E</a>: a neural network that creates images from text captions\n",
        "<li><a href='https://openai.com/blog/openai-codex/'>OpenAI Codex</a>: an AI system that translates natural language to code\n",
        "</ul>\n",
        "\n",
        "It was introduced in the seminal article <a href='https://arxiv.org/abs/1706.03762'>Attention Is All You Need</a> coauthored by Lukasz Kaiser.\n",
        "\n",
        "<p>We aim to explain the computer science theory behind this technology\n",
        "and to present it from basic principles without any use of AI jargon\n",
        "so that any mathematician can understand how it works and any programmer can implement it on their own.\n",
        "I studied the original article and the Trax implementation and produced this mathematical documentation.\n",
        "</p>\n",
        "\n",
        "<div>\n",
        "This article was written in 2021 by Michal Ryszard Wojcik (PhD in mathematics) under the guidance of Lukasz Kaiser (PhD in computer science) based on two sources:\n",
        "<ul><li>the article\n",
        "<a href='https://arxiv.org/abs/1706.03762'>Attention Is All You Need</a> coauthored by Lukasz Kaiser\n",
        "<li style='margin-top:0.5em;'>the source code\n",
        "<a href='https://github.com/google/trax/blob/master/trax/models/transformer.py#L194'>TransformerLM in Trax</a>\n",
        "supervised by Lukasz Kaiser\n",
        "</ul>\n",
        "</div>\n",
        "\n",
        "<p>Later I wrote a Python implementation of this mathematical specification, with some fully working demo examples:\n",
        "<a href='https://github.com/MichalRyszardWojcik/transformer-language-model/blob/main/TransformerLM_mathematical_definition_with_code_and_examples.ipynb'>github.com/MichalRyszardWojcik/transformer-language-model</a>.\n",
        "It can be executed in Google Colab directly in the browser without any installation.</p>\n",
        "\n",
        "<p>The official web page for this project is <a href='https://www.apronus.com/math/transformer-language-model-definition'>https://www.apronus.com/math/transformer-language-model-definition</a></p>"
      ],
      "metadata": {
        "id": "mK4UElDub4X-"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zyGB6jgqhFAM",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "bda83067-ecb7-4d9a-f6cf-93cee3a1257b"
      },
      "source": [
        "import numpy as np\n",
        "!pip install -q -U trax\n",
        "from trax import fastmath\n",
        "from trax.fastmath import numpy as jnp"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[K     |████████████████████████████████| 637 kB 31.6 MB/s \n",
            "\u001b[K     |████████████████████████████████| 4.6 MB 20.8 MB/s \n",
            "\u001b[K     |████████████████████████████████| 511.7 MB 6.4 kB/s \n",
            "\u001b[K     |████████████████████████████████| 5.8 MB 50.8 MB/s \n",
            "\u001b[K     |████████████████████████████████| 438 kB 71.8 MB/s \n",
            "\u001b[K     |████████████████████████████████| 1.6 MB 53.5 MB/s \n",
            "\u001b[?25h"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lAEGDiMrjwJp"
      },
      "source": [
        "https://github.com/google/trax/blob/master/trax/models/transformer.py#L182"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sl45bErH3zbP"
      },
      "source": [
        "# Rowwise and autoregressive matrix operations"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tee7DNXMrlLI"
      },
      "source": [
        "Let $T(X)=Y$ represent an operation which takes a matrix $X$ and returns a matrix $Y$ with the same number of rows as $X$, the number of colums being fixed for both input and output, say $m$ for input and $m'$ for output.\n",
        "\n",
        "We say that $T$ is **rowwise** iff for each matrix $X$ with $n$ rows and $m$ columns we have $T(X_i)=T(X)_i$ whenever $i\\leq n$, where $X_i$ is the i-th row of $X$ and $T(X)_i$ is the i-th row of $T(X)$.\n",
        "\n",
        "We say that $T$ is **autoregressive** iff for each matrix $X$ with $n$ rows and $m$ columns we have $T(X|i)=Y|i$ whenever $i\\leq n$, where the symbol $A|i$ denotes the matrix A reduced to the first $i$ rows, (e.g. $A|n=A$ and $A|1$ is the first row).\n",
        "\n",
        "Note that rowwise operations are autoregressive and compositions of autoregressive operations are autoregressive."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ez6iCdmVyIsm"
      },
      "source": [
        "Let $\\unicode{0x25FA}A$ denote the matrix $A$ modified so that it has $-\\infty$ above the main diagonal and is unchanged otherwise. Note that for a square matrix $A$ the rowwise operation softmax($\\unicode{0x25FA}A$) always returns a lower triangular square matrix."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hcXrkjWa0Fa2"
      },
      "source": [
        "**Claim.** If $A,B$ are autoregressive operations and $A$ always returns lower triangular square matrices then the matrix multiplication operation $X\\mapsto A(X)\\times B(X)$ is autoregressive."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5Bb6P0ORCRMM"
      },
      "source": [
        "Below we are going to define the Transformer Language Model and argue that it is autoregressive by appropriately applying the arguments collected above."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PWoKwbxXWHSA"
      },
      "source": [
        "def softmax_L(A):\n",
        "  X = jnp.tril(jnp.exp(A))\n",
        "  return X / jnp.sum(X,axis=1,keepdims=True)"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ejh9Eqgu_Bkp"
      },
      "source": [
        "# Embedding"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "clsp5gpT_Dyb"
      },
      "source": [
        "The Embedding layer is the first layer acting directly on the input composed of a sequence of tokens.\n",
        "\n",
        "Each token is an integer from $\\{0,1,\\ldots,\\mathrm{vocab\\_size}-1\\}$, where $\\mathrm{vocab\\_size}$ is one of the transformer language model's parameters.\n",
        "\n",
        "This layer translates tokens into \"vectorwords\" which are vectors in $\\mathbb R^{d_{model}}$, where $d_{model}$ is a parameter.\n",
        "\n",
        "There is an embedded array of weights of shape $\\mathrm{vocab\\_size}\\times d_{model}$. The layer simply assigns a row to each token.\n",
        "\n",
        "The output is a matrix of dimension $n\\times d_{model}$, where $n$ is the number of tokens in the input.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "x6Fw-ZEKVXDn"
      },
      "source": [
        "def embedding(weights, input):\n",
        "  n = input.shape[0]\n",
        "  rows = []\n",
        "  for i in range(n):\n",
        "    row = weights[input[i]]\n",
        "    rows.append( jnp.array([row]) )\n",
        "  return jnp.concatenate(tuple(rows), axis=0)"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_N8TRZEQvZPx"
      },
      "source": [
        "# Positional Encoding"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DSxGuvHZKL1J"
      },
      "source": [
        "Let us define the Positional Encoding layer in a transformer language model with max input length $M$ and $d_{model}$.\n",
        "\n",
        "Let $X$ be an input matrix with rows containing vectorwords of length $d_{model}$ with the number of rows between $1$ and $M$.\n",
        "\n",
        "Then $PE(X)=X+PE|n$ where $n$ is the number of rows in $X$ and $PE|n$ is the layer's embedded $PE$ weight matrix reduced to the first $n$ rows.\n",
        "\n",
        "Note that this layer simply adds its weight matrix to its input.\n",
        "\n",
        "Note that it is a rowwise operation."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ceQ8uLx06Jsq"
      },
      "source": [
        "def positionalEncoding(weights, X):\n",
        "  n = X.shape[0]\n",
        "  return X + weights[:n,:]"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9jJLgUPlS2B3"
      },
      "source": [
        "The weight matrix can be anything that results from the training process. But it is initialized in the following way."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3-owici1vc9x"
      },
      "source": [
        "$M$ is the maximum number of tokens in the input\n",
        "\n",
        "$$PE\\colon\\{1,\\ldots,M\\}\\times\\{1,\\ldots,d_{model}\\}\\to[-1,1]$$\n",
        "\n",
        "$$PE(pos,2i)=\\sin\\Big(\\cfrac{pos}{10000^{2i/d_{model}}}\\Big)$$\n",
        "\n",
        "$$PE(pos,2i+1)=\\cos\\Big(\\cfrac{pos}{10000^{2i/d_{model}}}\\Big)$$\n",
        "\n",
        "$PE$ is defined here as a function of two variables where the first variable is interpreted as the position in the input sequence (e.g. first token, second token, third token) and the second variable is interpreted as one of the $d_{model}$ axes of the vectorwords created from the input tokens.\n",
        "\n",
        "But we will also treat $PE$ as a $M\\times d_{model}$ matrix of weights embedded in the Positional Encoding layer.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kldYkvtD2v1P",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "30e268b0-0179-4854-e987-ff9233950a89"
      },
      "source": [
        "def PE(pos,k):\n",
        "  if k % 2 == 0:\n",
        "    trig = '\\mathrm{sin}'\n",
        "    i2 = str(k)\n",
        "  else:\n",
        "    trig = '\\mathrm{cos}'\n",
        "    i2 = str(k-1)\n",
        "  re = trig + \"\\Big(\\cfrac{\" + str(pos) + \"}{10000^{\" + i2 + \"/d_{model}}}\\Big)\"\n",
        "  return '$\\small{' + re + '}$ '\n",
        "\n",
        "def PEpos(pos,d_model):\n",
        "  ret = ''\n",
        "  for k in range(d_model):\n",
        "    ret += (PE(pos,k))\n",
        "  return ret + '\\n\\n'\n",
        "\n",
        "d_model = 6\n",
        "max_len = 3\n",
        "ret = ''\n",
        "for p in range(max_len):\n",
        "  ret += PEpos(p+1,d_model)\n",
        "print(ret)"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "$\\small{\\mathrm{sin}\\Big(\\cfrac{1}{10000^{0/d_{model}}}\\Big)}$ $\\small{\\mathrm{cos}\\Big(\\cfrac{1}{10000^{0/d_{model}}}\\Big)}$ $\\small{\\mathrm{sin}\\Big(\\cfrac{1}{10000^{2/d_{model}}}\\Big)}$ $\\small{\\mathrm{cos}\\Big(\\cfrac{1}{10000^{2/d_{model}}}\\Big)}$ $\\small{\\mathrm{sin}\\Big(\\cfrac{1}{10000^{4/d_{model}}}\\Big)}$ $\\small{\\mathrm{cos}\\Big(\\cfrac{1}{10000^{4/d_{model}}}\\Big)}$ \n",
            "\n",
            "$\\small{\\mathrm{sin}\\Big(\\cfrac{2}{10000^{0/d_{model}}}\\Big)}$ $\\small{\\mathrm{cos}\\Big(\\cfrac{2}{10000^{0/d_{model}}}\\Big)}$ $\\small{\\mathrm{sin}\\Big(\\cfrac{2}{10000^{2/d_{model}}}\\Big)}$ $\\small{\\mathrm{cos}\\Big(\\cfrac{2}{10000^{2/d_{model}}}\\Big)}$ $\\small{\\mathrm{sin}\\Big(\\cfrac{2}{10000^{4/d_{model}}}\\Big)}$ $\\small{\\mathrm{cos}\\Big(\\cfrac{2}{10000^{4/d_{model}}}\\Big)}$ \n",
            "\n",
            "$\\small{\\mathrm{sin}\\Big(\\cfrac{3}{10000^{0/d_{model}}}\\Big)}$ $\\small{\\mathrm{cos}\\Big(\\cfrac{3}{10000^{0/d_{model}}}\\Big)}$ $\\small{\\mathrm{sin}\\Big(\\cfrac{3}{10000^{2/d_{model}}}\\Big)}$ $\\small{\\mathrm{cos}\\Big(\\cfrac{3}{10000^{2/d_{model}}}\\Big)}$ $\\small{\\mathrm{sin}\\Big(\\cfrac{3}{10000^{4/d_{model}}}\\Big)}$ $\\small{\\mathrm{cos}\\Big(\\cfrac{3}{10000^{4/d_{model}}}\\Big)}$ \n",
            "\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5EAFXaxbFRjY"
      },
      "source": [
        "The image represents a vector of $d_{model}=6$ real numbers (vectorword).\n",
        "\n",
        "$\\Large[$$\\small{\\mathrm{sin}\\Big(\\cfrac{pos}{10000^{0/d_{model}}}\\Big)}$,\n",
        "$\\small{\\mathrm{cos}\\Big(\\cfrac{pos}{10000^{0/d_{model}}}\\Big)}$,\n",
        "$\\small{\\mathrm{sin}\\Big(\\cfrac{pos}{10000^{2/d_{model}}}\\Big)}$,\n",
        "$\\small{\\mathrm{cos}\\Big(\\cfrac{pos}{10000^{2/d_{model}}}\\Big)}$,\n",
        "$\\small{\\mathrm{sin}\\Big(\\cfrac{pos}{10000^{4/d_{model}}}\\Big)}$,\n",
        "$\\small{\\mathrm{cos}\\Big(\\cfrac{pos}{10000^{4/d_{model}}}\\Big)}$$\\Large]$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9Err41Z9JP6H"
      },
      "source": [
        "Let us assume maximum input length $M=3$ and $d_{model}=6$ for simplicity of presentation.\n",
        "\n",
        "The following $M\\times d_{model}$ matrix is populated with numbers from the $PE$ function."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h0EJCyMPDU7U"
      },
      "source": [
        "$\\small{\\mathrm{sin}\\Big(\\cfrac{1}{10000^{0/d_{model}}}\\Big)}$ $\\small{\\mathrm{cos}\\Big(\\cfrac{1}{10000^{0/d_{model}}}\\Big)}$ $\\small{\\mathrm{sin}\\Big(\\cfrac{1}{10000^{2/d_{model}}}\\Big)}$ $\\small{\\mathrm{cos}\\Big(\\cfrac{1}{10000^{2/d_{model}}}\\Big)}$ $\\small{\\mathrm{sin}\\Big(\\cfrac{1}{10000^{4/d_{model}}}\\Big)}$ $\\small{\\mathrm{cos}\\Big(\\cfrac{1}{10000^{4/d_{model}}}\\Big)}$ \n",
        "\n",
        "$\\small{\\mathrm{sin}\\Big(\\cfrac{2}{10000^{0/d_{model}}}\\Big)}$ $\\small{\\mathrm{cos}\\Big(\\cfrac{2}{10000^{0/d_{model}}}\\Big)}$ $\\small{\\mathrm{sin}\\Big(\\cfrac{2}{10000^{2/d_{model}}}\\Big)}$ $\\small{\\mathrm{cos}\\Big(\\cfrac{2}{10000^{2/d_{model}}}\\Big)}$ $\\small{\\mathrm{sin}\\Big(\\cfrac{2}{10000^{4/d_{model}}}\\Big)}$ $\\small{\\mathrm{cos}\\Big(\\cfrac{2}{10000^{4/d_{model}}}\\Big)}$ \n",
        "\n",
        "$\\small{\\mathrm{sin}\\Big(\\cfrac{3}{10000^{0/d_{model}}}\\Big)}$ $\\small{\\mathrm{cos}\\Big(\\cfrac{3}{10000^{0/d_{model}}}\\Big)}$ $\\small{\\mathrm{sin}\\Big(\\cfrac{3}{10000^{2/d_{model}}}\\Big)}$ $\\small{\\mathrm{cos}\\Big(\\cfrac{3}{10000^{2/d_{model}}}\\Big)}$ $\\small{\\mathrm{sin}\\Big(\\cfrac{3}{10000^{4/d_{model}}}\\Big)}$ $\\small{\\mathrm{cos}\\Big(\\cfrac{3}{10000^{4/d_{model}}}\\Big)}$ \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3PmBeofm5On0"
      },
      "source": [
        "From the paper *Attention Is All You Need*:\n",
        "We chose this function because we hypothesized it would allow the model to easily attend by relative position, since for any fixed offset $k$, $PE_{pos+k}$ can be represented as a linear function of $PE_{pos}$.\n",
        "\n",
        "The following articles help to understand the linear representation claim:\n",
        "\n",
        "https://kazemnejad.com/blog/transformer_architecture_positional_encoding/\n",
        "\n",
        "https://timodenk.com/blog/linear-relationships-in-the-transformers-positional-encoding/\n",
        "\n",
        "For any fixed offset $k$, there is a linear map $T_k\\colon\\mathbb R^{d_{model}}\\to\\mathbb R^{d_{model}}$ such that\n",
        "$$T_k\\circ PE_{pos}=PE_{pos+k}$$\n",
        "for any $pos\\in\\{1,2,\\ldots,M-k\\}$.\n",
        "\n",
        "For any fixed offset $k$, there is a linear map $T_k\\colon\\mathbb R^{d_{model}}\\to\\mathbb R^{d_{model}}$ such that\n",
        "$$T_k(PE(pos,i))=PE(pos+k,i)$$\n",
        "for any $pos\\in\\{1,2,\\ldots,M-k\\}$ and any $i\\in\\{1,\\ldots,d_{model}\\}$.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ilCp-wvx28pl"
      },
      "source": [
        "# Normalization layers"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1U_jZmdnt18u"
      },
      "source": [
        "Let us say that to standardize a vectorword $x\\in\\mathbb R^{d_{model}}$ means to subtract its arithmetic mean $\\overline{x}$ from each coordinate and divide each coordinate by $\\sigma(x)+\\epsilon$, where $\\sigma(x)$ is its standard deviation and $\\epsilon=10^{-6}$ is a constant embedded in the code.\n",
        "\n",
        "Let us say that to normalize a vectorword $x\\in\\mathbb R^{d_{model}}$ means to standardize it and then multiply it coordinatewise by a weight vector of the same length $a\\in\\mathbb R^{d_{model}}$ and then add a weight vector $b\\in\\mathbb R^{d_{model}}$.\n",
        "\n",
        "Each normalization layer has two weight vectors $a,b\\in\\mathbb R^{d_{model}}$.\n",
        "\n",
        "The input is as usual a matrix $X$ with $n$ rows, where $n$ is the number of words in the input sentence (the number of tokens in the input). Each row of $X$ is a vectorword from $\\mathbb R^{d_{model}}$.\n",
        "\n",
        "A normalization layer is a rowwise operation that normalizes each row using the same weights $a,b$.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Rcqwx-xS80QR"
      },
      "source": [
        "def normalize(weights,input,eps = 0.000001):\n",
        "  (a,b) = weights\n",
        "  mean = jnp.mean(input, axis=1, keepdims=True)\n",
        "  output = input - mean\n",
        "  std = jnp.std(output, axis=1, keepdims=True)\n",
        "  output = output / (std+eps)\n",
        "  return (output * a) + b"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZeojgnHI5cct"
      },
      "source": [
        "# DecoderBlock"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yqOcX1BvpbUs"
      },
      "source": [
        "**input**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sdqf-GB01Krw"
      },
      "source": [
        "X is a matrix with n rows, where n is the number of words in the input sentence (the number of tokens in the input)\n",
        "\n",
        "Each row of X is a vectorword from $R^{d_{model}}$ -- a vector of $d_{model}$ real numbers.\n",
        "\n",
        "Let us call this the usual shape for a matrix because such are the tensors accepted and returned by decoder blocks and its sublayers -- simplifying the exposition by assuming one sentence per batch and disregarding the batch size axis.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0KREvxF6nxBt"
      },
      "source": [
        "## Causal Attention with one head"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y6X85fHS-I1D"
      },
      "source": [
        "This section is not necessary because the Multi-Head Causal Attention section is self-contained and the definition of the Decoder Block refers to the Multi-Head Causal Attention section. The point is to be a gentle introduction to the multi-head version."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wnimzquiqLwr"
      },
      "source": [
        "**weights**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fDqSe7OOqsu0"
      },
      "source": [
        "$W^Q,W^K,W^V$ are square matrices of size ${d_{model}}\\times{d_{model}}$ (e.g. 512)\n",
        "\n",
        "They are the weights embedded in the Causal Attention layer CA."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1uKV488-qxc2"
      },
      "source": [
        "**definition**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Kln3ahDn3hWv"
      },
      "source": [
        "Let\n",
        "\n",
        "$Q(X)=X\\times W^Q$,\n",
        "\n",
        "$K(X)=X\\times W^K$,\n",
        "\n",
        "$V(X)=X\\times W^V$.\n",
        "\n",
        "Then the Causal Attention layer is defined as\n",
        "\n",
        "$$CA(X)=\\mathrm{softmax}\\Big(\\cfrac{\\unicode{0x25FA}\\big(Q(X)\\times K(X)^T\\big)}{\\sqrt{d_{model}}}\\Big)\\times V(X)$$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zGt0fC0b8wh_"
      },
      "source": [
        "Note that $CA$ is an autoregressive operation."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T1tnQ82f6rwI"
      },
      "source": [
        "## Multi-Head Causal Attention"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v9VWqz9p63jS"
      },
      "source": [
        "**parameters**\n",
        "\n",
        "$d_{model}$ is the length of the vectorwords\n",
        "\n",
        "$h$ is the number of attention heads\n",
        "\n",
        "$d_{head}=\\cfrac{d_{model}}{h}$ is the width of the weight matrices associated with individual heads"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F8dE-Fph6rwM"
      },
      "source": [
        "**weights**\n",
        "\n",
        "$W_i^Q,W_i^K,W_i^V$ are matrices of dimension ${d_{model}}\\times{d_{head}}$ for $i=1,\\ldots,h$\n",
        "\n",
        "$W^O$ is a matrix of dimension ${d_{model}}\\times{d_{model}}$\n",
        "\n",
        "$B$ is a matrix of dimension $n\\times{d_{model}}$ whose rows are identical"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NUQh6tSW6rwO"
      },
      "source": [
        "Fixing $i$, let\n",
        "\n",
        "$Q_i(X)=X\\times W_i^Q$,\n",
        "\n",
        "$K_i(X)=X\\times W_i^K$,\n",
        "\n",
        "$V_i(X)=X\\times W_i^V$.\n",
        "\n",
        "The dimension of these matrices is $n\\times d_{head}$.\n",
        "\n",
        "The calculation for an individual head is given by\n",
        "\n",
        "$$H_i(X)=\\mathrm{softmax}\\Big(\\unicode{0x25FA}\\Big(\\cfrac{Q_i(X)\\times K_i(X)^T}{\\sqrt{d_{head}}}\\Big)\\Big)\\times V_i(X)$$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "psBlseky6rwQ"
      },
      "source": [
        "Note that each $H_i$ is an autoregressive operation.\n",
        "\n",
        "Note that the output matrices $H_i(X)$ have dimension $n\\times d_{head}$.\n",
        "\n",
        "Let us concatenate them to obtain a matrix of dimension $n\\times d_{model}$\n",
        "\n",
        "$$H(X)=\\mathrm{Concat}(H_1,\\ldots,H_h)$$\n",
        "\n",
        "Note that $H$ is autoregressive."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jYy-df7RDPaD"
      },
      "source": [
        "Finally, the multi-head causal attention layer is defined by\n",
        "$$CA(X)=H(X)\\times W^O+B.$$\n",
        "Note that $CA$ is autoregressive."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "i9mK1TXwJLdI"
      },
      "source": [
        "def causalAttention(h,weights,X):\n",
        "  WQ, WK, WV, W0, B = weights\n",
        "  d_model = W0.shape[0]\n",
        "  d_head = d_model // h\n",
        "  denominator = jnp.sqrt(d_head)\n",
        "  def W_i(W,i):\n",
        "    c1 = d_head * i\n",
        "    c2 = d_head * (i+1)\n",
        "    return W[:,c1:c2]\n",
        "  def Q_i(i): return jnp.matmul(X,W_i(WQ,i))\n",
        "  def K_i(i): return jnp.matmul(X,W_i(WK,i))\n",
        "  def V_i(i): return jnp.matmul(X,W_i(WV,i))\n",
        "  def H_i(i):  \n",
        "    A = jnp.matmul(Q_i(i),jnp.transpose(K_i(i))) / denominator\n",
        "    return jnp.matmul(softmax_L(A),V_i(i))\n",
        "  H = jnp.concatenate(tuple( H_i(i) for i in range(h) ), axis=1)\n",
        "  return jnp.matmul(H,W0) + B"
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YropCDs8srMV"
      },
      "source": [
        "## The Feed Forward layer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wP45oy059wg2"
      },
      "source": [
        "The feed forward layer is defined by\n",
        "\n",
        "$$FF(X) = \\mathrm{ReLu}\\big(N_{ff}(X)\\times A+K\\big)\\times B+L$$\n",
        "\n",
        "where $N_{ff}$ is a normalization layer with its two weight vectors $a,b\\in\\mathbb R^{d_{model}}$,\n",
        "\n",
        "$A$ is a weight matrx of $d_{model}$ rows and $d_{ff}$ columns (e.g. $d_{ff} = 2048$),\n",
        "\n",
        "$K$ is a weight matrix of dimension $n\\times d_{ff}$ with identical rows,\n",
        "\n",
        "$B$ is a weight matrix of $d_{ff}$ rows and $d_{model}$ columns,\n",
        "\n",
        "$L$ is a weight matrix of dimension $n\\times d_{model}$ with identical rows.\n",
        "\n",
        "Note that $a,b,A,K,B,L$ are weights embedded in $FF$.\n",
        "\n",
        "Note that both the input and output are matrices of the usual shape of $n$ rows and $d_{model}$ columns."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S7AAovEmGMKX"
      },
      "source": [
        "Note that the input matrix $X$ and the output matrix $FF(X)$ have the same  usual shape $n\\times d_{model}$.\n",
        "\n",
        "Moreover, note that $FF$ is a rowwise operation.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1FYMHBhXBahI"
      },
      "source": [
        "def relu(X): return jnp.maximum(0,X)\n",
        "\n",
        "def feedforward(weights, X):\n",
        "  a, b, A, K, B, L = weights\n",
        "  q = jnp.matmul(normalize((a,b),X),A) + K\n",
        "  return jnp.matmul(relu(q),B) + L"
      ],
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aYpVe-ZC-4xu"
      },
      "source": [
        "## DecoderBlock Definition"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EZGZ-g0v_wZt"
      },
      "source": [
        "In order to match the Trax implementation we will write the definition using the residual layer notation $$Res(f)(x) = x+f(x),$$\n",
        "where $f$ is a layer and $x$ is an input tensor."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DPTVYHVi16JQ"
      },
      "source": [
        "The DecoderBlock is a composition of a normalization layer $N_{ca}(X)$ and the layers previously defined:\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AOyFFSbR-xRX"
      },
      "source": [
        "$$DB(X)=Res(FF)\\big(Res(CA\\circ N_{ca})(X)\\big)$$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "45VZOWd-AWzw"
      },
      "source": [
        "Without the residual layer notation it is simply\n",
        "\n",
        "$$DB(X)=X+CA(N_{ca}(X))+FF(X+CA(N_{ca}(X)))$$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o7Ct5iFoB-Gz"
      },
      "source": [
        "Note that the input matrix $X$ and the output matrix $DB(X)$ have the same  usual shape $n\\times d_{model}$.\n",
        "\n",
        "Moreover, note that $DB$ is an autoregressive operation."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ve-fdU7nvO5R",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "a0dd2ca6-44b3-4301-fb53-3d4579675354"
      },
      "source": [
        "def res(f,x): return f'{x}+{f}({x})'\n",
        "ca = res('CAoN','x')\n",
        "ff = res('FF',ca)\n",
        "print(ff)"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "x+CAoN(x)+FF(x+CAoN(x))\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bJpaCf0yA-VQ"
      },
      "source": [
        "**inspirational idea**\n",
        "\n",
        "Perhaps the decoder block would be a little more efficient if the formula was simplified to DB(X) = x+FF(x+CAoN(x))"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IHVmJfFCF8Lw"
      },
      "source": [
        "def decoderBlock(n_heads, weights, X):\n",
        "  Nweights, CAweights, FFweights = weights\n",
        "  CA = causalAttention(n_heads, CAweights, normalize(Nweights,X))\n",
        "  FF = feedforward(FFweights, X + CA)\n",
        "  return X + FF # according to the inspirational idea\n",
        "  return X + CA + FF # according to the definition"
      ],
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SGSYTbmPLO82"
      },
      "source": [
        "# Final Layer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cpY2nO8bLTc-"
      },
      "source": [
        "The input to the final layer is the output of a decoder block, which is a matrix $X$ of the usual shape $n\\times d_{model}$, where $n$ is the number of tokens in the original input to the transformer language model.\n",
        "\n",
        "The final layer's weight matrix $Y$ has dimension $d_{model}\\times\\mathrm{vocab\\_size}$.\n",
        "\n",
        "The final layer is a matrix multiplication defined as $\\Phi(X)=X\\times Y+B$,\n",
        "\n",
        "where $B$ is a weight matrix of dimension $n\\times\\mathrm{vocab\\_size}$ with identical rows.\n",
        "\n",
        "The output is a matrix of dimension $n\\times\\mathrm{vocab\\_size}$.\n",
        "\n",
        "Note that it is a rowwise operation."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LuSLEA8OJUKh"
      },
      "source": [
        "def finalLayer(weights, X):\n",
        "  Y, B = weights\n",
        "  return jnp.matmul(X,Y) + B"
      ],
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9XcrmS9lEgh8"
      },
      "source": [
        "# Definition of TransformerLM"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a37rqcP1Owzo"
      },
      "source": [
        "**parameters**\n",
        "\n",
        "$\\mathrm{vocab\\_size}$ is the size of the set of input tokens\n",
        "\n",
        "$d_{model}=512$ is the length of the vectorwords passed from sublayer to sublayer\n",
        "\n",
        "$d_{ff}=2048$ is the width of certain weight matrices inside the feed forward layer in each decoder block\n",
        "\n",
        "$n\\_layers=6$ is the number of decoder blocks\n",
        "\n",
        "$n\\_heads=8$ is the number of attention heads in each decoder block\n",
        "\n",
        "$\\mathrm{max\\_len}=2048$ is the maximum number of tokens in the input sequence\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U2PYbXspElKz"
      },
      "source": [
        "$$T_{LM}=\\Phi\\circ N\\circ DB_{n\\_layers}\\circ\\ldots\\circ DB_1\\circ PE\\circ Em$$\n",
        "\n",
        "where $N$ is a normalization layer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1aHwwCzfOI2f"
      },
      "source": [
        "The input is a sequence of tokens.\n",
        "\n",
        "Each token is an integer from $\\{0,1,\\ldots,\\mathrm{vocab\\_size}-1\\}$.\n",
        "\n",
        "The output is a matrix of dimension $n\\times\\mathrm{vocab\\_size}$ containing real numbers, where $n$ is the number of tokens in the input sequence.\n",
        "\n",
        "Note that it is an autoregressive operation if we interpret the input as a matrix of dimension $n\\times 1$."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "T9k9yWFEJ3j_"
      },
      "source": [
        "def transformerLM(n_heads, weights, input):\n",
        "  EMweights, PEweights, DBweights, Nweights, Fweights = weights\n",
        "  EM = embedding(EMweights, input)\n",
        "  PE = positionalEncoding(PEweights, EM)\n",
        "  n_layers = len(DBweights)\n",
        "  X = PE\n",
        "  for i in range(n_layers):\n",
        "    X = decoderBlock(n_heads, DBweights[i], X)\n",
        "  X = normalize(Nweights, X)\n",
        "  return finalLayer(Fweights, X)"
      ],
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DEBG8u3ePFiA"
      },
      "source": [
        "def generateRandomWeights(vocab_size, d_model = 512, d_ff = 2048, n_layers = 6, n_heads = 8, max_len = 2048, real = float):\n",
        "  rng = np.random.default_rng()\n",
        "  def randomweights(shape):\n",
        "    return (rng.random(shape) - 0.5).astype(real)\n",
        "  def generateRandomDBweights():\n",
        "    Nweights = (randomweights((d_model,)), randomweights((d_model,)))\n",
        "    shape = (d_model, d_model)\n",
        "    WQ = randomweights(shape)\n",
        "    WK = randomweights(shape)\n",
        "    WV = randomweights(shape)\n",
        "    W0 = randomweights(shape)\n",
        "    B = randomweights((1,d_model))\n",
        "    CAweights = (WQ, WK, WV, W0, B)\n",
        "    a, b = randomweights((d_model,)), randomweights((d_model,))\n",
        "    A = randomweights((d_model, d_ff))\n",
        "    K = randomweights((1,d_ff))\n",
        "    B = randomweights((d_ff,d_model))\n",
        "    L = randomweights((1,d_model))\n",
        "    FFweights = (a, b, A, K, B, L)\n",
        "    return (Nweights, CAweights, FFweights)\n",
        "  EMshape = (vocab_size, d_model)\n",
        "  EMweights = randomweights(EMshape)\n",
        "  PEshape = (max_len, d_model)\n",
        "  PEweights = randomweights(PEshape)\n",
        "  DBweights = []\n",
        "  for i in range(n_layers):\n",
        "    DBweights.append(generateRandomDBweights())\n",
        "  DBweights = tuple(DBweights)\n",
        "  Nshape = (d_model,)\n",
        "  Nweights = (randomweights(Nshape),randomweights(Nshape))\n",
        "  Yshape = (d_model, vocab_size)\n",
        "  Bshape = (1, vocab_size)\n",
        "  Fweights = (randomweights(Yshape), randomweights(Bshape))\n",
        "\n",
        "  weights = (EMweights, PEweights, DBweights, Nweights, Fweights)\n",
        "  return weights"
      ],
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sFoQa4oogdkI"
      },
      "source": [
        "def generateRandomInput(vocab_size, n):\n",
        "  rng = np.random.default_rng()\n",
        "  return rng.integers(0,vocab_size,n)"
      ],
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "buaDX44uegIF"
      },
      "source": [
        "# Evaluation (the loss function)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m8GvUA5ap8LY"
      },
      "source": [
        "## The essence of the loss function &mdash; single sequence input with no loss weights"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L7KnAxJrejBs"
      },
      "source": [
        "The input sequences are elements of $\\{0,1,\\ldots,\\rm{vocab\\_size-1}\\}^n$, where $n$ is the sequence length and $\\rm{vocab\\_size}$ is the number of possible tokens from which the sequences are built.\n",
        "\n",
        "We are given an autoregressive operation $T$, which takes an input sequence (interpreted as a matrix of dimension $n\\times 1$) and outputs a real-valued matrix of dimension $n\\times\\rm{vocab\\_size}$.\n",
        "\n",
        "In this context, we are going to define a loss function which returns a positive real number for each input sequence.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VMN-cxsRjR1N"
      },
      "source": [
        "Let $(x_1,\\ldots,x_n)$ be the input sequence.\n",
        "\n",
        "In the first step of the definition,\n",
        "let $X=(0,x_1,\\ldots,x_n)$ be viewed as the $(n+1)\\times 1$ dimensional input matrix for the autoregressive operation $T$. Then $T(X)$ is a real-valued matrix of dimension $(n+1)\\times\\rm{vocab\\_size}$, but we will be interested only in the first $n$ rows $T(X)|n$.\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZDJ6Vrj-mwQz"
      },
      "source": [
        "Note that each row of $T(X)|n$ is a sequence of $\\rm{vocab\\_size}$ many arbitrary real numbers as outputted by the working of the operation $T$. In the second step, we apply the LogSoftmax function to each row separately to obtain $Y=\\rm{LogSoftmax}(T(X)|n)$, which is now a matrix of $n$ rows with each row serving as a log probability distribution over the set of tokens."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8VrdqdUEpyAh"
      },
      "source": [
        "At this point we have created a function $k\\mapsto Y[k]$, which assigns a log probability distribution over the set of tokens for each position in the original input sequence $(x_1,\\ldots,x_n)$.\n",
        "\n",
        "Note that --- due to the autoregressive nature of $T$ and the insertion of the $0$ token in the first step --- we have guaranteed that:\n",
        "* the log distribution $Y[1]$ does not depend on the input sequence at all\n",
        "* the log distribution $Y[2]$ depends only on $x_1$\n",
        "* the log distribution $Y[3]$ is a function of $(x_1,x_2)$\n",
        "* the log distribution $Y[k]$ is a function of $(x_1,x_2,\\ldots,x_{k-1})$ for each $k=4,5,\\ldots,n$.\n",
        "\n",
        "Therefore it makes sense to evaluate the outputted log distribution $Y[k]$ against the actual token $x_k$ in the input."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xpwLw70XtqQq"
      },
      "source": [
        "Let $y_k = Y[k]_i$, where $i=x_k+1$, for each $k=1,\\ldots,n$.\n",
        "\n",
        "Note that in our setup $y_k\\lt 0$ is interpreted as the log probability of the token $x_k$ appearing on the $k$th position in the input string after $(x_1,\\ldots,x_{k-1})$. The closer to zero the better the evaluation."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w4-hkflezfdy"
      },
      "source": [
        "Finally, the loss function is defined as\n",
        "$$-\\cfrac{1}{n}\\sum_{i=1}^n y_k>0.$$\n",
        "The closer to zero the better."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jlOOaALneZQM"
      },
      "source": [
        "## Batch evaluation with loss weights"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NL-rlVzIeg_B"
      },
      "source": [
        "The input sequences are elements of $\\{0,1,\\ldots,\\rm{vocab\\_size-1}\\}^n$, where $n$ is the sequence length and $\\rm{vocab\\_size}$ is the number of possible tokens from which the sequences are built.\n",
        "\n",
        "We will be dealing with a batch of input sequences &mdash; formally a matrix of dimension $\\rm{batch\\_size}\\times n$, whose rows are input sequences. Additionally, there's a matrix of the same dimension whose elements are real numbers from the closed unit interval $[0,1]$, which is called the matrix of loss weights for the input batch.\n",
        "\n",
        "We are given an autoregressive operation $T$, which takes an input sequence (interpreted as a matrix of dimension $n\\times 1$) and outputs a real-valued matrix of dimension $n\\times\\rm{vocab\\_size}$. This time the operation $T$ is going to act simultaneously and independently on each input sequence from the batch so that formally it is a function which takes an array of dimension $\\rm{batch\\_size}\\times n$ and returns an array of dimension $\\rm{batch\\_size}\\times n\\times\\rm{vocab\\_size}$.\n",
        "\n",
        "In this context, we are going to define a loss function which returns a positive real number for each batch of input sequences with loss weights.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6UU5Zr9ciEQe"
      },
      "source": [
        "Let us fix notation so that the $j$th row in the input batch is denoted as $(x_{j1},\\ldots,x_{jn})$ for $j=1,\\ldots,\\rm{batch\\_size}$.\n",
        "\n",
        "Similarly, let the $j$th row in the loss weights matrix be denoted as $(\\omega_{j1},\\ldots,\\omega_{jn})$ for $j=1,\\ldots,\\rm{batch\\_size}$."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-gFBMNI1pJRe"
      },
      "source": [
        "Let $(y_{j1},\\ldots,y_{jn})$ be the log probabilities for the $j$th input sequence as computed according to the formula in the previous section, for $j=1,\\ldots,\\rm{batch\\_size}$."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fL-OdY2urvhR"
      },
      "source": [
        "The loss function is defined by\n",
        "$$-\\cfrac{\\sum_{(j,i)}\\omega_{ji}y_{ji}}{\\sum_{(j,i)}\\omega_{ji}}$$\n",
        "where the summation index $(j,i)$ ranges over\n",
        "$\\{1,\\ldots,\\rm{batch\\_size}\\}\\times\\{1,\\ldots,n\\}$.\n",
        "\n",
        "Note that this is a generalization of the loss function from the previous section, which can be seen by setting the weights to zeros on all rows except on a single row where they should be all ones."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UitH_YdEEH5L"
      },
      "source": [
        "### not the same as average over rows"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MkWc0wzNxJSP"
      },
      "source": [
        "Note that this is <b>not the same</b> thing as the average loss function of a single input sequence with loss weights taken over all the rows:\n",
        "$$\\cfrac{1}{s}\\sum_{j=1}^s\\Big(-\\cfrac{\\sum_{i=1}^n\\omega_{ji}y_{ji}}{\\sum_{i=1}^n\\omega_{ji}}\\Big)$$\n",
        "The two formulas give the same result if the sums $\\sum_{i=1}^n\\omega_{ji}$ are the same for each row, but they are slightly different otherwise. The following code illustrates these points."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "t69tslgh69yE",
        "outputId": "93934180-1930-41f3-83e1-44eebe443325"
      },
      "source": [
        "import numpy as np\n",
        "s = np.random.randint(1,89)\n",
        "n = np.random.randint(1,86)\n",
        "y = np.random.random_sample(size=(s,n))\n",
        "w = np.random.random_sample(size=(s,n))\n",
        "def formula1(y,w):\n",
        "  a = 0.0\n",
        "  b = 0.0\n",
        "  for j in range(s):\n",
        "    for i in range(n):\n",
        "      a += w[j][i]*y[j][i]\n",
        "      b += w[j][i]\n",
        "  return a/b\n",
        "def formula2(y,w):\n",
        "  c = 0.0\n",
        "  for j in range(s):\n",
        "    a = 0.0\n",
        "    b = 0.0\n",
        "    for i in range(n):\n",
        "      a += w[j][i]*y[j][i]\n",
        "      b += w[j][i]\n",
        "    c += a/b\n",
        "  return c/s\n",
        "print(formula1(y,w))\n",
        "print(formula2(y,w))\n",
        "# expected to be slightly different"
      ],
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0.5006191316067764\n",
            "0.5015580101356863\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "c5tOj_3N-_g1",
        "outputId": "fe87e019-bfd5-4f1d-9d19-8413bede7053"
      },
      "source": [
        "import numpy as np\n",
        "s = np.random.randint(1,89)\n",
        "n = np.random.randint(1,86)\n",
        "_y = np.random.random_sample(size=(s,n))\n",
        "_w = np.random.random_sample(size=(1,n))\n",
        "def w(j,i):\n",
        "  return _w[0][(i+j)%n]\n",
        "  return _w[0][i]\n",
        "def y(j,i):\n",
        "  return _y[j][i]\n",
        "def formula1():\n",
        "  a = 0.0\n",
        "  b = 0.0\n",
        "  for j in range(s):\n",
        "    for i in range(n):\n",
        "      a += w(j,i)*y(j,i)\n",
        "      b += w(j,i)\n",
        "  return a/b\n",
        "def formula2():\n",
        "  c = 0.0\n",
        "  for j in range(s):\n",
        "    a = 0.0\n",
        "    b = 0.0\n",
        "    for i in range(n):\n",
        "      a += w(j,i)*y(j,i)\n",
        "      b += w(j,i)\n",
        "    c += a/b\n",
        "  return c/s\n",
        "print(formula1())\n",
        "print(formula2())\n",
        "# expected to be equal"
      ],
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0.48544311762896336\n",
            "0.48544311762896436\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nXD1JV0JHsdX"
      },
      "source": [
        "## The code for the loss function"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ettcuhp-VgCa"
      },
      "source": [
        "def LogSoftmax(x):\n",
        "  return x - fastmath.logsumexp(x, -1, keepdims=True)"
      ],
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def weighted_loss(T, input, mask):\n",
        "  '''\n",
        "  T is a function\n",
        "         with input shape (batch_size, n)\n",
        "          and output shape (batch_size, n, vocab_size)\n",
        "  input, mask both have shape (batch_size, n)\n",
        "  '''\n",
        "  batch_size = input.shape[0]\n",
        "  n = input.shape[1] # the length of input sequences\n",
        "  zero = jnp.zeros((batch_size,1), np.int32)\n",
        "  X = jnp.concatenate((zero,input), axis=1)\n",
        "  # X.shape == (batch_size, n+1)\n",
        "  Y = LogSoftmax( T(X)[:n,:] )\n",
        "  # Y.shape[0] == batch_size\n",
        "  # Y.shape[1] == n+1\n",
        "  # Y.shape[2] == vocab_size\n",
        "  y = []\n",
        "  for row_number in range(batch_size):\n",
        "    row = []\n",
        "    for token_position in range(n):\n",
        "      logprob = Y[row_number, token_position, input[row_number, token_position]]\n",
        "      row.append(logprob)\n",
        "    y.append(row)\n",
        "  y = jnp.array(y)\n",
        "  # y.shape == mask.shape\n",
        "  return - jnp.sum(mask*y) / jnp.sum(mask)"
      ],
      "metadata": {
        "id": "fdQmQu9J15f_"
      },
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WJrLbFIGTo59"
      },
      "source": [
        "## Gradient-Descent-based minimization of the loss function"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ODpCHGPMTslC"
      },
      "source": [
        "Our loss function is a positive real-valued function but we need to take a closer look at its domain and its differentiability in order to use it within a gradient-descent-based minimization algorithm."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N-V81uf8VYE0"
      },
      "source": [
        "The loss function depends on the batch of input sequences on the one hand and on the collection of weights from the $T$ operation (which is a TransformerLM). We may think of a batch space and a weight space so that the loss function takes a point from the batch space and a point from the weight space to return the loss value.\n",
        "\n",
        "Note that the loss function is suitable for a mini-batch (stochastic) gradient descent setup because it works for any point from the batch space.\n",
        "\n",
        "<b>Claim.</b>\n",
        "Fixing a batch space point, this function is differentiable almost everywhere as a multivariate real-valued function from the weight space.\n",
        "<br><i>Proof.</i>\n",
        "Apart from ReLu it is a composition of the basic arithmetic operations (addition, subtraction, multiplication, division) and the exp and log elementary functions from softmax and logSoftmax.\n",
        "\n",
        "In order to calculate the gradient by automatic differentiation it remains to arbitrarily define the derivative of ReLu at 0.\n",
        "\n",
        "(There is also the square root in the calculation of the std.)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Vulh4PyEcQW8"
      },
      "source": [
        "Useful links for this subsection:\n",
        "* [Gradient Descent For Machine Learning\n",
        "by Jason Brownlee](https://machinelearningmastery.com/gradient-descent-for-machine-learning/)\n",
        "* [Batch, Mini Batch & Stochastic Gradient Descent by Sushant Patrikar](https://towardsdatascience.com/batch-mini-batch-stochastic-gradient-descent-7a62ecba642a)\n",
        "* [An overview of gradient descent optimization algorithms by Sebastian Ruder](https://ruder.io/optimizing-gradient-descent/)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lNXxc6PHwols"
      },
      "source": [
        "# Fully working demo"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UXB2K1eMKedg"
      },
      "source": [
        "## The dataset for demo purposes"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "y7k9g0xNKiP5"
      },
      "source": [
        "import random\n",
        "\n",
        "def generate_batch(batch_size, n_tokens, n):\n",
        "  shape = (batch_size, 2*n+2)\n",
        "  a = np.zeros(shape,np.int32)\n",
        "  for y in range(batch_size):\n",
        "    for i in range(n):\n",
        "      x = random.randint(1,n_tokens)\n",
        "      a[y][i] = x\n",
        "      a[y][2*n-i] = x\n",
        "  return a\n",
        "\n",
        "def train_stream(batch_size,\n",
        "                 n_tokens,\n",
        "                 input_length_min,\n",
        "                 input_length_max):\n",
        "  n = random.randint(input_length_min, input_length_max)\n",
        "  batch = generate_batch(batch_size, n_tokens, n)\n",
        "  inputs = batch\n",
        "  shape = (batch_size, 2*n+2)\n",
        "  loss_weights = np.zeros(shape,np.float32)\n",
        "  for y in range(batch_size):\n",
        "    for i in range(n+1,2*n+2):\n",
        "      loss_weights[y][i] = 1.0\n",
        "  return (inputs,loss_weights)\n",
        "\n",
        "def data_batches_yielder(batch_size, input_params):\n",
        "  n_tokens, input_length_min, input_length_max = input_params\n",
        "  while True:\n",
        "    input, loss_weights = train_stream(batch_size, n_tokens, input_length_min, input_length_max)\n",
        "    yield (input, loss_weights)"
      ],
      "execution_count": 19,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "shyqj1GCitXj"
      },
      "source": [
        "## The training process"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SXhE9CIfKX3H"
      },
      "source": [
        "def initWeights(model_params):\n",
        "  (vocab_size, d_model, d_ff, n_layers, n_heads, max_len) = model_params\n",
        "  real = float\n",
        "  return generateRandomWeights(vocab_size, d_model, d_ff, n_layers, n_heads, max_len, real)"
      ],
      "execution_count": 20,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hxJACg6Pn8mB"
      },
      "source": [
        "import time\n",
        "\n",
        "def train(batch_size, n_steps, input_params, model_params, weights):\n",
        "  start_time = time.time()\n",
        "  data_batches = data_batches_yielder(batch_size, input_params)\n",
        "  n_heads = model_params[4]\n",
        "\n",
        "  def network(data_batch, weights):\n",
        "    def T(X):\n",
        "      rows = []\n",
        "      for j in range(batch_size):\n",
        "        row = transformerLM(n_heads, weights, X[j])\n",
        "        rows.append(row)\n",
        "      return jnp.array(rows)\n",
        "    input, mask = data_batch\n",
        "    return weighted_loss(T, input, mask)\n",
        "\n",
        "  value_and_grad = fastmath.value_and_grad(network, argnums=1)\n",
        "  value_and_grad = fastmath.jit(value_and_grad)\n",
        "\n",
        "  print(f'Starting the {n_steps}-steps training process...')\n",
        "  acc_loss = 0.0\n",
        "  for i in range(n_steps):\n",
        "    data = next(data_batches)\n",
        "    loss, grad = value_and_grad(data, weights)\n",
        "    acc_loss += loss\n",
        "    weights = fastmath.nested_map_multiarg(lambda w, g: w - 0.001 * g, weights, grad)\n",
        "    if i % 100 == 99:\n",
        "      print(f'avg loss at {i+1} steps: {acc_loss / 100:.3f}')\n",
        "      acc_loss = 0.0\n",
        "  end_time = time.time()\n",
        "  seconds = int(end_time - start_time)\n",
        "  minutes = int(seconds / 60)\n",
        "  print(f'The training process took {minutes}:{seconds % 60}.')\n",
        "  return weights"
      ],
      "execution_count": 21,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nklV_w5SD8rc"
      },
      "source": [
        "## Testing the trained model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_mcAA-dKN5sK"
      },
      "source": [
        "def transformerLM_output(model_params, weights, input, length=0):\n",
        "  (vocab_size, d_model, d_ff, n_layers, n_heads, max_len) = model_params\n",
        "  if length == 0: length = max_len\n",
        "  X = [0] + list(input)\n",
        "  while len(X) < length:\n",
        "    Y = transformerLM(n_heads, weights, np.array(X))[-1,:]\n",
        "    token = int(np.argmax(Y))\n",
        "    X.append(token)\n",
        "    if token == 0: break\n",
        "  return X[1:]\n",
        "\n",
        "'''\n",
        "def onesequencetest(input_params, model_params, weights):\n",
        "  n_tokens, input_length_min, input_length_max = input_params\n",
        "  n = random.randint(input_length_min, input_length_max)\n",
        "  perfect = generate_batch(1, n_tokens, n)[0]\n",
        "  input = perfect[:n+1]\n",
        "  n_heads = model_params[4]\n",
        "  output = transformerLM_output(model_params, weights, input)\n",
        "  if (output == list(perfect)):\n",
        "    print('correct')\n",
        "  else:\n",
        "    print('wrong')\n",
        "  print(list(perfect))\n",
        "  print(output)\n",
        "  return 1 if (output == list(perfect)) else 0\n",
        "'''\n",
        "\n",
        "def onesequencetest(input_params, model_params, weights):\n",
        "  n_tokens, input_length_min, input_length_max = input_params\n",
        "  n = random.randint(input_length_min, input_length_max)\n",
        "  perfect = generate_batch(1, n_tokens, n)[0]\n",
        "  input = perfect[:n+1]\n",
        "  n_heads = model_params[4]\n",
        "  output = transformerLM_output(model_params, weights, input)\n",
        "  if (output != list(perfect)):\n",
        "    print('wrong')\n",
        "    print('should be:', list(perfect))\n",
        "    print('instead:  ', output)\n",
        "    return 0\n",
        "  return 1\n",
        "\n",
        "def runtest(input_params, model_params, weights, n):\n",
        "  print(f'Running a test on {n} random sequences...')\n",
        "  r = []\n",
        "  for i in range(n):\n",
        "    r.append(onesequencetest(input_params, model_params, weights))\n",
        "  print(100*sum(r)/n, 'success rate')"
      ],
      "execution_count": 22,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Lh9bdNIln1Wr"
      },
      "source": [
        "## Demo examples"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GoL44ZzV4oKD"
      },
      "source": [
        "In these examples the task is to reverse the order of tokens in the input sequence."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DymumQAR4Tix"
      },
      "source": [
        "### Example 0"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "U0eLe7Df4Scy"
      },
      "source": [
        "def example0():\n",
        "  n_tokens = 10\n",
        "  min_length = 2\n",
        "  max_length = 2\n",
        "  # 100 elements in the training dataset\n",
        "  input_params = (n_tokens, min_length, max_length)\n",
        "\n",
        "  vocab_size = n_tokens + 1\n",
        "  d_model = 128\n",
        "  d_ff = 256\n",
        "  n_layers = 2\n",
        "  n_heads = 2\n",
        "  max_len = 2 * max_length + 2 + 1\n",
        "  model_params = (vocab_size, d_model, d_ff, n_layers, n_heads, max_len)\n",
        "\n",
        "  weights = initWeights(model_params)\n",
        "  batch_size = 4\n",
        "  n_steps = 6000\n",
        "  weights = train(batch_size,n_steps,input_params, model_params, weights)\n",
        "\n",
        "  runtest(input_params, model_params, weights, 100)"
      ],
      "execution_count": 23,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# example0() # less than two minutes"
      ],
      "metadata": {
        "id": "yfPpxqTRGses"
      },
      "execution_count": 24,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MKqUrpkkpg6x"
      },
      "source": [
        "### Example 1"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1vm-Z5nvn4qx"
      },
      "source": [
        "def example1():\n",
        "  n_tokens = 3\n",
        "  min_length = 2\n",
        "  max_length = 3\n",
        "  # 36 elements in the training dataset\n",
        "  input_params = (n_tokens, min_length, max_length)\n",
        "\n",
        "  vocab_size = n_tokens + 1\n",
        "  d_model = 128\n",
        "  d_ff = 256\n",
        "  n_layers = 2\n",
        "  n_heads = 2\n",
        "  max_len = 2 * max_length + 2 + 1\n",
        "  model_params = (vocab_size, d_model, d_ff, n_layers, n_heads, max_len)\n",
        "\n",
        "  weights = initWeights(model_params)\n",
        "  batch_size = 4\n",
        "  n_steps = 6000\n",
        "  weights = train(batch_size, n_steps, input_params, model_params, weights)\n",
        "\n",
        "  runtest(input_params, model_params, weights, 100)"
      ],
      "execution_count": 25,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# example1() # less than 2 minutes"
      ],
      "metadata": {
        "id": "JgMd0dcMMvaV"
      },
      "execution_count": 26,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yujPjvUdvg61"
      },
      "source": [
        "### Example 2"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AEo2YsRivjEt"
      },
      "source": [
        "def example2():\n",
        "  n_tokens = 4\n",
        "  min_length = 2\n",
        "  max_length = 4\n",
        "  # 336 elements in the training dataset\n",
        "  input_params = (n_tokens, min_length, max_length)\n",
        "\n",
        "  vocab_size = n_tokens + 1\n",
        "  d_model = 128\n",
        "  d_ff = 256\n",
        "  n_layers = 2\n",
        "  n_heads = 2\n",
        "  max_len = 2 * max_length + 2 + 1\n",
        "  model_params = (vocab_size, d_model, d_ff, n_layers, n_heads, max_len)\n",
        "\n",
        "  weights = initWeights(model_params)\n",
        "\n",
        "  batch_size = 4\n",
        "  n_steps = 16_000\n",
        "  weights = train(batch_size, n_steps, input_params, model_params, weights)\n",
        "\n",
        "  runtest(input_params, model_params, weights, 100)"
      ],
      "execution_count": 27,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# example2() # less than 5 minutes"
      ],
      "metadata": {
        "id": "65vW2iZ2Obds"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}