llm-self-attention.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "inputs = torch.tensor(\n",
    "  [[0.43, 0.15, 0.89], # Your     (x^1)\n",
    "   [0.55, 0.87, 0.66], # journey  (x^2)\n",
    "   [0.57, 0.85, 0.64], # starts   (x^3)\n",
    "   [0.22, 0.58, 0.33], # with     (x^4)\n",
    "   [0.77, 0.25, 0.10], # one      (x^5)\n",
    "   [0.05, 0.80, 0.55]] # step     (x^6)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865])\n"
     ]
    }
   ],
   "source": [
    "query = inputs[1]\n",
    "attn_scores_2 = torch.empty(inputs.shape[0])\n",
    "for i, x_i in enumerate(inputs):\n",
    "    attn_scores_2[i] = torch.dot(x_i, query)\n",
    "print(attn_scores_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.9544)\n",
      "tensor(0.9544)\n"
     ]
    }
   ],
   "source": [
    "res = 0.\n",
    "for idx, element in enumerate(inputs[0]):\n",
    "    res += inputs[0][idx] * query[idx]\n",
    "print(res)\n",
    "print(torch.dot(inputs[0], query))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Attention weights: tensor([0.1455, 0.2278, 0.2249, 0.1285, 0.1077, 0.1656])\n",
      "Sum: tensor(1.0000)\n"
     ]
    }
   ],
   "source": [
    "attn_weights_2_tmp = attn_scores_2 / attn_scores_2.sum()\n",
    "print(\"Attention weights:\", attn_weights_2_tmp)\n",
    "print(\"Sum:\", attn_weights_2_tmp.sum())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Attention weights: tensor([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581])\n",
      "Sum: tensor(1.)\n"
     ]
    }
   ],
   "source": [
    "def softmax_naive(x):\n",
    "    return torch.exp(x) / torch.exp(x).sum(dim=0)\n",
    " \n",
    "attn_weights_2_naive = softmax_naive(attn_scores_2)\n",
    "print(\"Attention weights:\", attn_weights_2_naive)\n",
    "print(\"Sum:\", attn_weights_2_naive.sum())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Attention weights: tensor([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581])\n",
      "Sum: tensor(1.)\n"
     ]
    }
   ],
   "source": [
    "attn_weights_2 = torch.softmax(attn_scores_2, dim=0)\n",
    "print(\"Attention weights:\", attn_weights_2)\n",
    "print(\"Sum:\", attn_weights_2.sum())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.4419, 0.6515, 0.5683])\n"
     ]
    }
   ],
   "source": [
    "query = inputs[1] # 2nd input token is the query\n",
    "context_vec_2 = torch.zeros(query.shape)\n",
    "for i,x_i in enumerate(inputs):\n",
    "    context_vec_2 += attn_weights_2[i]*x_i\n",
    "print(context_vec_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.9995, 0.9544, 0.9422, 0.4753, 0.4576, 0.6310],\n",
      "        [0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865],\n",
      "        [0.9422, 1.4754, 1.4570, 0.8296, 0.7154, 1.0605],\n",
      "        [0.4753, 0.8434, 0.8296, 0.4937, 0.3474, 0.6565],\n",
      "        [0.4576, 0.7070, 0.7154, 0.3474, 0.6654, 0.2935],\n",
      "        [0.6310, 1.0865, 1.0605, 0.6565, 0.2935, 0.9450]])\n"
     ]
    }
   ],
   "source": [
    "attn_scores = torch.empty(6, 6)\n",
    "for i, x_i in enumerate(inputs):\n",
    "    for j, x_j in enumerate(inputs):\n",
    "        attn_scores[i, j] = torch.dot(x_i, x_j)\n",
    "print(attn_scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.9995, 0.9544, 0.9422, 0.4753, 0.4576, 0.6310],\n",
      "        [0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865],\n",
      "        [0.9422, 1.4754, 1.4570, 0.8296, 0.7154, 1.0605],\n",
      "        [0.4753, 0.8434, 0.8296, 0.4937, 0.3474, 0.6565],\n",
      "        [0.4576, 0.7070, 0.7154, 0.3474, 0.6654, 0.2935],\n",
      "        [0.6310, 1.0865, 1.0605, 0.6565, 0.2935, 0.9450]])\n"
     ]
    }
   ],
   "source": [
    "attn_scores = inputs @ inputs.T\n",
    "print(attn_scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.2098, 0.2006, 0.1981, 0.1242, 0.1220, 0.1452],\n",
      "        [0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581],\n",
      "        [0.1390, 0.2369, 0.2326, 0.1242, 0.1108, 0.1565],\n",
      "        [0.1435, 0.2074, 0.2046, 0.1462, 0.1263, 0.1720],\n",
      "        [0.1526, 0.1958, 0.1975, 0.1367, 0.1879, 0.1295],\n",
      "        [0.1385, 0.2184, 0.2128, 0.1420, 0.0988, 0.1896]])\n"
     ]
    }
   ],
   "source": [
    "attn_weights = torch.softmax(attn_scores, dim=-1)\n",
    "print(attn_weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Row 2 sum: 1.0\n",
      "All row sums: tensor([1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000])\n"
     ]
    }
   ],
   "source": [
    "row_2_sum = sum([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581])\n",
    "print(\"Row 2 sum:\", row_2_sum)\n",
    "print(\"All row sums:\", attn_weights.sum(dim=-1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.4421, 0.5931, 0.5790],\n",
      "        [0.4419, 0.6515, 0.5683],\n",
      "        [0.4431, 0.6496, 0.5671],\n",
      "        [0.4304, 0.6298, 0.5510],\n",
      "        [0.4671, 0.5910, 0.5266],\n",
      "        [0.4177, 0.6503, 0.5645]])\n"
     ]
    }
   ],
   "source": [
    "all_context_vecs = attn_weights @ inputs\n",
    "print(all_context_vecs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Previous 2nd context vector: tensor([0.4419, 0.6515, 0.5683])\n"
     ]
    }
   ],
   "source": [
    "print(\"Previous 2nd context vector:\", context_vec_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_2 = inputs[1]\n",
    "d_in = inputs.shape[1]\n",
    "d_out = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(123)\n",
    "W_query = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)\n",
    "W_key   = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)\n",
    "W_value = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.4306, 1.4551])\n"
     ]
    }
   ],
   "source": [
    "query_2 = x_2 @ W_query \n",
    "key_2 = x_2 @ W_key \n",
    "value_2 = x_2 @ W_value\n",
    "print(query_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "keys.shape: torch.Size([6, 2])\n",
      "values.shape: torch.Size([6, 2])\n"
     ]
    }
   ],
   "source": [
    "keys = inputs @ W_key \n",
    "values = inputs @ W_value\n",
    "print(\"keys.shape:\", keys.shape)\n",
    "print(\"values.shape:\", values.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "class SelfAttention_v1(nn.Module):\n",
    "    def __init__(self, d_in, d_out):\n",
    "        super().__init__()\n",
    "        self.W_query = nn.Parameter(torch.rand(d_in, d_out))\n",
    "        self.W_key   = nn.Parameter(torch.rand(d_in, d_out))\n",
    "        self.W_value = nn.Parameter(torch.rand(d_in, d_out))\n",
    " \n",
    "    def forward(self, x):\n",
    "        keys = x @ self.W_key\n",
    "        queries = x @ self.W_query\n",
    "        values = x @ self.W_value\n",
    "        attn_scores = queries @ keys.T # omega\n",
    "        attn_weights = torch.softmax(\n",
    "            attn_scores / keys.shape[-1]**0.5, dim=-1)\n",
    "        context_vec = attn_weights @ values\n",
    "        return context_vec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.2996, 0.8053],\n",
      "        [0.3061, 0.8210],\n",
      "        [0.3058, 0.8203],\n",
      "        [0.2948, 0.7939],\n",
      "        [0.2927, 0.7891],\n",
      "        [0.2990, 0.8040]], grad_fn=<MmBackward0>)\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(123)\n",
    "sa_v1 = SelfAttention_v1(d_in, d_out)\n",
    "print(sa_v1(inputs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SelfAttention_v2(nn.Module):\n",
    "    def __init__(self, d_in, d_out, qkv_bias=False):\n",
    "        super().__init__()\n",
    "        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.W_key   = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    "        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
    " \n",
    "    def forward(self, x):\n",
    "        keys = self.W_key(x)\n",
    "        queries = self.W_query(x)\n",
    "        values = self.W_value(x)\n",
    "        attn_scores = queries @ keys.T\n",
    "        attn_weights = torch.softmax(\n",
    "            attn_scores / keys.shape[-1]**0.5, dim=-1\n",
    "        )\n",
    "        context_vec = attn_weights @ values\n",
    "        return context_vec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-0.0739,  0.0713],\n",
      "        [-0.0748,  0.0703],\n",
      "        [-0.0749,  0.0702],\n",
      "        [-0.0760,  0.0685],\n",
      "        [-0.0763,  0.0679],\n",
      "        [-0.0754,  0.0693]], grad_fn=<MmBackward0>)\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(789)\n",
    "sa_v2 = SelfAttention_v2(d_in, d_out)\n",
    "print(sa_v2(inputs))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}