multi-headed_self-attention.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multi-Headed Self-Attention\n",
    "![Multi-Head Attention](./images/multi-head-attention.png)\n",
    "\n",
    "We will build an autoregressive multi-headed self-attention language model for predicting the next token in a sequence of text. The shakespeare_char dataset will be used for this demonstration, which can be found in the data folder.  \n",
    "\n",
    "A multi-headed self-attention block combines multiple smaller single-headed self-attention blocks and concatenates the output of all of them. It can extract information from different representation subspaces. This enables it to capture more diverse and complex patterns in the input space, while single-headed self-attention can only focus on a single subspace at a time. \n",
    "\n",
    "We use a head size of the single-heads that is equal to the number of embedding dimensions divided by the number of heads. This ensures that the output dimensions of the multi-headed block is the same as a single-headed block, provided that 'num_heads' can divide 'n_embed' without any remainder. For further explanation on multi-headed attention see Andrej Karpathy's video [2].\n",
    "\n",
    "\n",
    "### References:\n",
    "- [1] [GPT colab notebook](https://colab.research.google.com/drive/1JMLa53HDuA-i7ZBmqV7ZnA3c_fvtXnx-?usp=sharing)\n",
    "- [2] [Video: multi-headed self-attention](https://www.youtube.com/watch?v=kCc8FmEb1nY&t=4919s)\n",
    "- [3] [Attention Is All You Need](https://arxiv.org/abs/1706.03762)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import requests\n",
    "import numpy as np\n",
    "import jax\n",
    "from jax import random\n",
    "import jax.numpy as jnp\n",
    "import flax.linen as nn\n",
    "import optax\n",
    "from jax import value_and_grad\n",
    "\n",
    "from helper_funcs import get_batch, generate, masked_fill, loss_fn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_embed = 32 # Number of embedding dimensions\n",
    "batch_size = 32 # How many independent sequences will we process in parallel?\n",
    "block_size = 8 # What is the maximum context length for predictions?\n",
    "num_heads = 4 # Number of heads in the multi-headed block\n",
    "\n",
    "rng_key = jax.random.PRNGKey(128)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "length of dataset in characters: 1,115,394\n",
      "all the unique characters: \n",
      " !$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz\n",
      "vocab size: 65\n",
      "train has 1,003,854 tokens\n",
      "val has 111,540 tokens\n"
     ]
    }
   ],
   "source": [
    "# download the tiny shakespeare dataset\n",
    "input_file_path = os.path.join('./data/shakespeare_char/input.txt')\n",
    "if not os.path.exists(input_file_path):\n",
    "    data_url = 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'\n",
    "    with open(input_file_path, 'w') as f:\n",
    "        f.write(requests.get(data_url).text)\n",
    "\n",
    "with open(input_file_path, 'r') as f:\n",
    "    data = f.read()\n",
    "print(f\"length of dataset in characters: {len(data):,}\")\n",
    "\n",
    "# get all the unique characters that occur in this text\n",
    "chars = sorted(list(set(data)))\n",
    "vocab_size = len(chars)\n",
    "print(\"all the unique characters:\", ''.join(chars))\n",
    "print(f\"vocab size: {vocab_size:,}\")\n",
    "\n",
    "# create a mapping from characters to integers\n",
    "stoi = { ch:i for i,ch in enumerate(chars) }\n",
    "itos = { i:ch for i,ch in enumerate(chars) }\n",
    "def encode(s):\n",
    "    return [stoi[c] for c in s] # encoder: take a string, output a list of integers\n",
    "def decode(l):\n",
    "    l = np.array(l)\n",
    "    return ''.join([itos[i] for i in l]) # decoder: take a list of integers, output a string\n",
    "\n",
    "# create the train and test splits\n",
    "n = len(data)\n",
    "train_data = data[:int(n*0.9)]\n",
    "val_data = data[int(n*0.9):]\n",
    "\n",
    "# encode both to integers\n",
    "train_ids = encode(train_data)\n",
    "val_ids = encode(val_data)\n",
    "print(f\"train has {len(train_ids):,} tokens\")\n",
    "print(f\"val has {len(val_ids):,} tokens\")\n",
    "\n",
    "train_ids = jnp.array(train_ids, dtype=jnp.uint16)\n",
    "val_ids = jnp.array(val_ids, dtype=jnp.uint16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First Citizen:\n",
      "Before we proceed any further, hear me speak.\n",
      "\n",
      "All:\n",
      "Speak, speak.\n",
      "\n",
      "First Citizen:\n",
      "You\n"
     ]
    }
   ],
   "source": [
    "print(decode(train_ids[:100]))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build the Attention Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Head(nn.Module):\n",
    "    \"\"\"\n",
    "    A single-headed self-attention decoder block.\n",
    "    Takes the combined token and position embedding as input,\n",
    "    then calculates the key and query values.\n",
    "    The key and query are multiplied to calculate the \n",
    "    attention scores/affinities. The future weights are\n",
    "    then altered to have zero affinity, this ensures the \n",
    "    model can't \"cheat\". The input is then used to calculate\n",
    "    the values, which are then aggregated by multiplying \n",
    "    them with the weights.\n",
    "    \"\"\"\n",
    "    head_size: int\n",
    "\n",
    "    @nn.compact\n",
    "    def __call__(self, x):\n",
    "        B,T,C = x.shape\n",
    "        key = nn.Dense(self.head_size, use_bias=False)\n",
    "        k = key(x) # (B,T,C)\n",
    "        query = nn.Dense(self.head_size, use_bias=False)\n",
    "        q = query(x) # (B,T,C)\n",
    "        # compute attention scores (\"affinities\")\n",
    "        weights =  q @ k.transpose((0, -1, -2)) * self.head_size**-0.5 # (B, T, C) @ (B, C, T) ---> (B, T, T)\n",
    "        tril = jnp.tril(jnp.ones(shape=(T, T), dtype=bool))\n",
    "        tril = jnp.repeat(tril[None, ...], repeats=B, axis=0)\n",
    "        weights = masked_fill(tril, weights, -jnp.inf)\n",
    "        weights = jax.nn.softmax(weights, axis=-1)\n",
    "        # perform the weighted aggregation of the values\n",
    "        value = nn.Dense(self.head_size, use_bias=False)\n",
    "        v = value(x)\n",
    "        out = weights @ v\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadedAttention(nn.Module):\n",
    "    \"\"\"\n",
    "    Combines multiple heads of scaled self-attention \n",
    "    in parallel, then concatenates the heads outputs.\n",
    "    \"\"\"\n",
    "    num_heads: int\n",
    "    head_size: int\n",
    "\n",
    "    @nn.compact\n",
    "    def __call__(self, x):\n",
    "        # Create a list of num_heads heads\n",
    "        heads = [Head(self.head_size) for _ in range(self.num_heads)]\n",
    "        # Provide the same input for each head\n",
    "        heads_out = [h(x) for h in heads]\n",
    "        combined_logits = jnp.concatenate(heads_out, axis=-1)\n",
    "        return combined_logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AttentionLanguageModel(nn.Module):\n",
    "    \"\"\"\n",
    "    Multi-headed self-attention language model.\n",
    "    Uses the previous tokens in the sequence to \n",
    "    determine the probabilities of the next token.\n",
    "    Processes the combined position and token embedding\n",
    "    through a multi-headed self-attention decoder block, \n",
    "    which is then processed through a dense layer to \n",
    "    aquire the token logits.\n",
    "    The logits can then be processed through a softmax\n",
    "    function to calculate the token probabilities.\n",
    "    \"\"\"\n",
    "    vocab_size: int\n",
    "    n_embed: int\n",
    "    block_size: int\n",
    "    num_heads: int\n",
    "    \n",
    "    @nn.compact\n",
    "    def __call__(self, index_seq):\n",
    "        B, T = index_seq.shape\n",
    "\n",
    "        # Each token directly reads off the logits for the next token from a lookup table\n",
    "        token_embedding_table = nn.Embed(num_embeddings=self.vocab_size, features=self.n_embed) \n",
    "        token_emb = token_embedding_table(index_seq) # (B, T, C)\n",
    "\n",
    "        position_embedding_table = nn.Embed(num_embeddings=self.block_size, features=self.n_embed) \n",
    "        pos_emb = position_embedding_table(jnp.arange(T)) # (T, C)\n",
    "\n",
    "        x = token_emb + pos_emb # (B, T, C)\n",
    "\n",
    "        head_size = self.n_embed // self.num_heads\n",
    "        sa_heads = MultiHeadedAttention(self.num_heads, head_size)\n",
    "        x = sa_heads(x) # apply one head of self-attention (B, T, C)\n",
    "\n",
    "        lm_head = nn.Dense(self.vocab_size)\n",
    "        logits = lm_head(x) # (B, T, vocab_size)\n",
    "\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = AttentionLanguageModel(vocab_size, n_embed, block_size, num_heads)\n",
    "dummy_x = jnp.zeros(shape=(batch_size, block_size), dtype=jnp.uint16)\n",
    "variables = model.init(rng_key, dummy_x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(32, 8, 65)\n"
     ]
    }
   ],
   "source": [
    "out = model.apply(variables, dummy_x)\n",
    "print(out.shape)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Text Generation Pre-Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated text: \n",
      "\n",
      "FeRkiTwf.'jtMwetQ\n",
      "y;:zYFfVmFiPuyZaYrv,wzii Tgh,j.rE.'srlHn&QEy'sjb33e&?K3,FEgIeNBOm zv;vZlQlNm.mqbbM\n"
     ]
    }
   ],
   "source": [
    "index_seq = jnp.zeros(shape=(1,1), dtype=jnp.uint16)\n",
    "max_new_tokens = 100\n",
    "\n",
    "generated_indices = generate(variables, model.apply, index_seq, rng_key, vocab_size, 1, block_size, max_new_tokens)\n",
    "generated_indices = list(np.array(generated_indices[0]))\n",
    "print(\"Generated text: \")\n",
    "print(decode(generated_indices))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = optax.adamw(learning_rate=1e-2)\n",
    "opt_state = optimizer.init(variables)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0, Loss: 4.2002\n",
      "Epoch: 1, Loss: 4.1314\n",
      "Epoch: 2, Loss: 4.0487\n",
      "Epoch: 3, Loss: 3.9995\n",
      "Epoch: 4, Loss: 3.9011\n",
      "Epoch: 5, Loss: 3.7680\n",
      "Epoch: 6, Loss: 3.6813\n",
      "Epoch: 7, Loss: 3.5244\n",
      "Epoch: 8, Loss: 3.4426\n",
      "Epoch: 9, Loss: 3.4472\n",
      "Epoch: 10, Loss: 3.4264\n",
      "Epoch: 11, Loss: 3.4981\n",
      "Epoch: 12, Loss: 3.1786\n",
      "Epoch: 13, Loss: 3.2236\n",
      "Epoch: 14, Loss: 3.6526\n",
      "Epoch: 15, Loss: 3.4811\n",
      "Epoch: 16, Loss: 3.3162\n",
      "Epoch: 17, Loss: 3.2343\n",
      "Epoch: 18, Loss: 3.2692\n",
      "Epoch: 19, Loss: 3.3367\n",
      "Epoch: 20, Loss: 3.1827\n",
      "Epoch: 21, Loss: 3.2085\n",
      "Epoch: 22, Loss: 3.1524\n",
      "Epoch: 23, Loss: 3.1159\n",
      "Epoch: 24, Loss: 3.1055\n",
      "Epoch: 25, Loss: 3.1955\n",
      "Epoch: 26, Loss: 3.2321\n",
      "Epoch: 27, Loss: 3.0995\n",
      "Epoch: 28, Loss: 3.1132\n",
      "Epoch: 29, Loss: 3.1211\n",
      "Epoch: 30, Loss: 3.1240\n",
      "Epoch: 31, Loss: 3.1205\n",
      "Epoch: 32, Loss: 3.0943\n",
      "Epoch: 33, Loss: 3.0640\n",
      "Epoch: 34, Loss: 3.0053\n",
      "Epoch: 35, Loss: 3.0887\n",
      "Epoch: 36, Loss: 2.9509\n",
      "Epoch: 37, Loss: 2.9637\n",
      "Epoch: 38, Loss: 2.8833\n",
      "Epoch: 39, Loss: 3.0016\n",
      "Epoch: 40, Loss: 2.8503\n",
      "Epoch: 41, Loss: 3.0264\n",
      "Epoch: 42, Loss: 2.9108\n",
      "Epoch: 43, Loss: 2.9234\n",
      "Epoch: 44, Loss: 2.9406\n",
      "Epoch: 45, Loss: 3.1706\n",
      "Epoch: 46, Loss: 3.0657\n",
      "Epoch: 47, Loss: 2.7789\n",
      "Epoch: 48, Loss: 2.7642\n",
      "Epoch: 49, Loss: 3.0235\n",
      "Epoch: 50, Loss: 2.7490\n",
      "Epoch: 51, Loss: 2.7665\n",
      "Epoch: 52, Loss: 2.8009\n",
      "Epoch: 53, Loss: 2.6438\n",
      "Epoch: 54, Loss: 2.7681\n",
      "Epoch: 55, Loss: 2.8252\n",
      "Epoch: 56, Loss: 2.7645\n",
      "Epoch: 57, Loss: 2.8408\n",
      "Epoch: 58, Loss: 2.8685\n",
      "Epoch: 59, Loss: 2.8447\n",
      "Epoch: 60, Loss: 2.7370\n",
      "Epoch: 61, Loss: 2.6097\n",
      "Epoch: 62, Loss: 2.7552\n",
      "Epoch: 63, Loss: 2.7740\n",
      "Epoch: 64, Loss: 2.8124\n",
      "Epoch: 65, Loss: 2.6498\n",
      "Epoch: 66, Loss: 2.8426\n",
      "Epoch: 67, Loss: 2.6394\n",
      "Epoch: 68, Loss: 2.7828\n",
      "Epoch: 69, Loss: 2.6226\n",
      "Epoch: 70, Loss: 2.5779\n",
      "Epoch: 71, Loss: 2.5664\n",
      "Epoch: 72, Loss: 2.6307\n",
      "Epoch: 73, Loss: 2.5260\n",
      "Epoch: 74, Loss: 2.7697\n",
      "Epoch: 75, Loss: 2.6794\n",
      "Epoch: 76, Loss: 2.5365\n",
      "Epoch: 77, Loss: 2.8592\n",
      "Epoch: 78, Loss: 2.5805\n",
      "Epoch: 79, Loss: 2.7155\n",
      "Epoch: 80, Loss: 2.4368\n",
      "Epoch: 81, Loss: 2.5010\n",
      "Epoch: 82, Loss: 2.5617\n",
      "Epoch: 83, Loss: 2.8457\n",
      "Epoch: 84, Loss: 2.7391\n",
      "Epoch: 85, Loss: 2.5087\n",
      "Epoch: 86, Loss: 2.6226\n",
      "Epoch: 87, Loss: 2.6016\n",
      "Epoch: 88, Loss: 2.6035\n",
      "Epoch: 89, Loss: 2.8334\n",
      "Epoch: 90, Loss: 2.5495\n",
      "Epoch: 91, Loss: 2.7333\n",
      "Epoch: 92, Loss: 2.6847\n",
      "Epoch: 93, Loss: 2.5214\n",
      "Epoch: 94, Loss: 2.5889\n",
      "Epoch: 95, Loss: 2.6230\n",
      "Epoch: 96, Loss: 2.6512\n",
      "Epoch: 97, Loss: 2.4776\n",
      "Epoch: 98, Loss: 2.6964\n",
      "Epoch: 99, Loss: 2.6593\n"
     ]
    }
   ],
   "source": [
    "steps = 100\n",
    "\n",
    "for step in range(steps):\n",
    "    rng_key, subkey = jax.random.split(rng_key)\n",
    "    xb, yb = get_batch(train_ids, subkey, batch_size, block_size)\n",
    "\n",
    "    loss, grads = value_and_grad(loss_fn, argnums=(0))(\n",
    "        variables, \n",
    "        model.apply,\n",
    "        xb, \n",
    "        yb\n",
    "    )\n",
    "    updates, opt_state = optimizer.update(grads, opt_state, variables)\n",
    "    variables = optax.apply_updates(variables, updates) \n",
    "\n",
    "    print(f\"Epoch: {step}, Loss: {loss :.4f}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Text Generation Post-Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated text: \n",
      "\n",
      "Tot hte; hef preer tusl ous blrlee.\n",
      "\n",
      "N Gowrvee.\n",
      "\n",
      "\n",
      "S yorc.\n",
      "\n",
      "o\n",
      "L:\n",
      "Fe've,\n",
      ":\n",
      "\n",
      "Tetscshe Yfodot d muj inOt\n"
     ]
    }
   ],
   "source": [
    "index_seq = jnp.zeros(shape=(1,1), dtype=jnp.uint16)\n",
    "max_new_tokens = 100\n",
    "\n",
    "generated_indices = generate(variables, model.apply, index_seq, rng_key, vocab_size, 1, block_size, max_new_tokens)\n",
    "generated_indices = list(np.array(generated_indices[0]))\n",
    "print(\"Generated text: \")\n",
    "print(decode(generated_indices))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "flax2",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "61a659281f728d55fc033691f61c254a0324a8156d8c726d9ae327f7f6499663"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}