4_overfitting_single_batch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Contents\n",
    "1. Generate a single toy batch from the tiny shakespeare dataset\n",
    "1. Overfit on that batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from gpt2 import GPT, GPTConfig # our GPT class\n",
    "import tiktoken\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "enc = tiktoken.get_encoding('gpt2')\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # dynamic device"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating a batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Length of tokens: 297884\n",
      "input: \n",
      "tensor([[ 5962, 22307,    25,  8421,   356,  5120],\n",
      "        [  597,  2252,    11,  3285,   502,  2740],\n",
      "        [   13,  3237,    25,  5248,   461,    11],\n",
      "        [ 2740,    13,  5962, 22307,    25,  1639]])\n",
      "target: \n",
      "tensor([[22307,    25,  8421,   356,  5120,   597],\n",
      "        [ 2252,    11,  3285,   502,  2740,    13],\n",
      "        [ 3237,    25,  5248,   461,    11,  2740],\n",
      "        [   13,  5962, 22307,    25,  1639,   389]])\n"
     ]
    }
   ],
   "source": [
    "# read input data text file\n",
    "with open('data/input.txt', 'r') as file:\n",
    "    data = file.read().replace('\\n', '')\n",
    "\n",
    "tokens = enc.encode(data)\n",
    "print(f'Length of tokens: {len(tokens)}')\n",
    "\n",
    "B, T = 4, 6\n",
    "x = torch.tensor(tokens[:B*T]).reshape(B, T)\n",
    "y = torch.tensor(tokens[1:B*T+1]).reshape(B, T)\n",
    "print(f'input: \\n{x}')\n",
    "print(f'target: \\n{y}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overfitting on a single batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = GPT(GPTConfig).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loss at iteration 0: 11.043309211730957\n",
      "Loss at iteration 1: 8.613590240478516\n",
      "Loss at iteration 2: 10.750606536865234\n",
      "Loss at iteration 3: 8.073920249938965\n",
      "Loss at iteration 4: 7.72338342666626\n",
      "Loss at iteration 5: 7.411949157714844\n",
      "Loss at iteration 6: 6.878050327301025\n",
      "Loss at iteration 7: 6.415774345397949\n",
      "Loss at iteration 8: 5.869007110595703\n",
      "Loss at iteration 9: 5.304551601409912\n",
      "Loss at iteration 10: 4.673495769500732\n",
      "Loss at iteration 11: 4.278286457061768\n",
      "Loss at iteration 12: 4.2021050453186035\n",
      "Loss at iteration 13: 3.4737212657928467\n",
      "Loss at iteration 14: 3.127682685852051\n",
      "Loss at iteration 15: 2.6242685317993164\n",
      "Loss at iteration 16: 2.2600083351135254\n",
      "Loss at iteration 17: 1.910429835319519\n",
      "Loss at iteration 18: 1.6248867511749268\n",
      "Loss at iteration 19: 1.3340888023376465\n",
      "Loss at iteration 20: 1.0824331045150757\n",
      "Loss at iteration 21: 0.8567084074020386\n",
      "Loss at iteration 22: 0.6526219844818115\n",
      "Loss at iteration 23: 0.4932633936405182\n",
      "Loss at iteration 24: 0.38657405972480774\n",
      "Loss at iteration 25: 0.31045642495155334\n",
      "Loss at iteration 26: 0.24792629480361938\n",
      "Loss at iteration 27: 0.19973434507846832\n",
      "Loss at iteration 28: 0.16331036388874054\n",
      "Loss at iteration 29: 0.13431665301322937\n",
      "Loss at iteration 30: 0.11090229451656342\n",
      "Loss at iteration 31: 0.09239556640386581\n",
      "Loss at iteration 32: 0.07773861289024353\n",
      "Loss at iteration 33: 0.06571812927722931\n",
      "Loss at iteration 34: 0.055486053228378296\n",
      "Loss at iteration 35: 0.04679794982075691\n",
      "Loss at iteration 36: 0.03971758857369423\n",
      "Loss at iteration 37: 0.03411246836185455\n",
      "Loss at iteration 38: 0.029604041948914528\n",
      "Loss at iteration 39: 0.025846343487501144\n",
      "Loss at iteration 40: 0.022678261622786522\n",
      "Loss at iteration 41: 0.020047426223754883\n",
      "Loss at iteration 42: 0.017894674092531204\n",
      "Loss at iteration 43: 0.016130205243825912\n",
      "Loss at iteration 44: 0.014665314927697182\n",
      "Loss at iteration 45: 0.013431834056973457\n",
      "Loss at iteration 46: 0.012381810694932938\n",
      "Loss at iteration 47: 0.011479987762868404\n",
      "Loss at iteration 48: 0.010698860511183739\n",
      "Loss at iteration 49: 0.010016019456088543\n",
      "Loss at iteration 50: 0.009413325227797031\n",
      "Loss at iteration 51: 0.008876289241015911\n",
      "Loss at iteration 52: 0.008393746800720692\n",
      "Loss at iteration 53: 0.00795731134712696\n",
      "Loss at iteration 54: 0.0075605944730341434\n",
      "Loss at iteration 55: 0.007198920473456383\n",
      "Loss at iteration 56: 0.006868466734886169\n",
      "Loss at iteration 57: 0.006566220428794622\n",
      "Loss at iteration 58: 0.006289579905569553\n",
      "Loss at iteration 59: 0.0060361819341778755\n",
      "Loss at iteration 60: 0.0058037531562149525\n",
      "Loss at iteration 61: 0.005590352229773998\n",
      "Loss at iteration 62: 0.005394033156335354\n",
      "Loss at iteration 63: 0.005213071592152119\n",
      "Loss at iteration 64: 0.005045877303928137\n",
      "Loss at iteration 65: 0.004891033750027418\n",
      "Loss at iteration 66: 0.004747205879539251\n",
      "Loss at iteration 67: 0.004613411612808704\n",
      "Loss at iteration 68: 0.0044886087998747826\n",
      "Loss at iteration 69: 0.004371990915387869\n",
      "Loss at iteration 70: 0.004262847360223532\n",
      "Loss at iteration 71: 0.004160548094660044\n",
      "Loss at iteration 72: 0.004064497072249651\n",
      "Loss at iteration 73: 0.003974144347012043\n",
      "Loss at iteration 74: 0.0038891416043043137\n",
      "Loss at iteration 75: 0.003809055546298623\n",
      "Loss at iteration 76: 0.0037334461230784655\n",
      "Loss at iteration 77: 0.003661961993202567\n",
      "Loss at iteration 78: 0.0035943591501563787\n",
      "Loss at iteration 79: 0.003530277870595455\n",
      "Loss at iteration 80: 0.003469446673989296\n",
      "Loss at iteration 81: 0.0034116017632186413\n",
      "Loss at iteration 82: 0.0033565741032361984\n",
      "Loss at iteration 83: 0.003304105019196868\n",
      "Loss at iteration 84: 0.0032539418898522854\n",
      "Loss at iteration 85: 0.0032059710938483477\n",
      "Loss at iteration 86: 0.0031600692309439182\n",
      "Loss at iteration 87: 0.0031160491053014994\n",
      "Loss at iteration 88: 0.003073745174333453\n",
      "Loss at iteration 89: 0.003033030079677701\n",
      "Loss at iteration 90: 0.002993872156366706\n",
      "Loss at iteration 91: 0.002956127980723977\n",
      "Loss at iteration 92: 0.002919713268056512\n",
      "Loss at iteration 93: 0.002884554211050272\n",
      "Loss at iteration 94: 0.0028505309019237757\n",
      "Loss at iteration 95: 0.0028176503255963326\n",
      "Loss at iteration 96: 0.002785810735076666\n",
      "Loss at iteration 97: 0.0027549401856958866\n",
      "Loss at iteration 98: 0.0027249655686318874\n",
      "Loss at iteration 99: 0.0026959015522152185\n"
     ]
    }
   ],
   "source": [
    "B, T = 4, 32\n",
    "x = torch.tensor(tokens[:B*T]).reshape(B, T).to(device)\n",
    "y = torch.tensor(tokens[1:B*T+1]).reshape(B, T).to(device)\n",
    "\n",
    "model.train();\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)\n",
    "for i in range(100):\n",
    "    optimizer.zero_grad()\n",
    "    logits, loss = model(x, y)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    print(f'Loss at iteration {i}: {loss.item()}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Why does not tying the weights make the loss only decrease up to around `1.0`? Removing the weight tying allows the loss to go well below `1.0`."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}