self_attention.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x: torch.Size([4, 8, 2])\n",
      "---- b: 0\n",
      "-- t: 0, cur_and_prev: torch.Size([1, 2]), mean: tensor([ 0.1808, -0.0700])\n",
      "-- t: 1, cur_and_prev: torch.Size([2, 2]), mean: tensor([-0.0894, -0.4926])\n",
      "-- t: 2, cur_and_prev: torch.Size([3, 2]), mean: tensor([ 0.1490, -0.3199])\n",
      "-- t: 3, cur_and_prev: torch.Size([4, 2]), mean: tensor([ 0.3504, -0.2238])\n",
      "-- t: 4, cur_and_prev: torch.Size([5, 2]), mean: tensor([0.3525, 0.0545])\n",
      "-- t: 5, cur_and_prev: torch.Size([6, 2]), mean: tensor([ 0.0688, -0.0396])\n",
      "-- t: 6, cur_and_prev: torch.Size([7, 2]), mean: tensor([ 0.0927, -0.0682])\n",
      "-- t: 7, cur_and_prev: torch.Size([8, 2]), mean: tensor([-0.0341,  0.1332])\n",
      "---- b: 1\n",
      "-- t: 0, cur_and_prev: torch.Size([1, 2]), mean: tensor([ 1.3488, -0.1396])\n",
      "-- t: 1, cur_and_prev: torch.Size([2, 2]), mean: tensor([0.8173, 0.4127])\n",
      "-- t: 2, cur_and_prev: torch.Size([3, 2]), mean: tensor([-0.1342,  0.4395])\n",
      "-- t: 3, cur_and_prev: torch.Size([4, 2]), mean: tensor([0.2711, 0.4774])\n",
      "-- t: 4, cur_and_prev: torch.Size([5, 2]), mean: tensor([0.2421, 0.0694])\n",
      "-- t: 5, cur_and_prev: torch.Size([6, 2]), mean: tensor([0.0084, 0.0020])\n",
      "-- t: 6, cur_and_prev: torch.Size([7, 2]), mean: tensor([ 0.0712, -0.1128])\n",
      "-- t: 7, cur_and_prev: torch.Size([8, 2]), mean: tensor([0.2527, 0.2149])\n",
      "---- b: 2\n",
      "-- t: 0, cur_and_prev: torch.Size([1, 2]), mean: tensor([-0.6631, -0.2513])\n",
      "-- t: 1, cur_and_prev: torch.Size([2, 2]), mean: tensor([ 0.1735, -0.0649])\n",
      "-- t: 2, cur_and_prev: torch.Size([3, 2]), mean: tensor([0.1685, 0.3348])\n",
      "-- t: 3, cur_and_prev: torch.Size([4, 2]), mean: tensor([-0.1621,  0.1765])\n",
      "-- t: 4, cur_and_prev: torch.Size([5, 2]), mean: tensor([-0.2312, -0.0436])\n",
      "-- t: 5, cur_and_prev: torch.Size([6, 2]), mean: tensor([-0.1015, -0.2855])\n",
      "-- t: 6, cur_and_prev: torch.Size([7, 2]), mean: tensor([-0.2593, -0.1630])\n",
      "-- t: 7, cur_and_prev: torch.Size([8, 2]), mean: tensor([-0.3015, -0.2293])\n",
      "---- b: 3\n",
      "-- t: 0, cur_and_prev: torch.Size([1, 2]), mean: tensor([ 1.6455, -0.8030])\n",
      "-- t: 1, cur_and_prev: torch.Size([2, 2]), mean: tensor([ 1.4985, -0.5395])\n",
      "-- t: 2, cur_and_prev: torch.Size([3, 2]), mean: tensor([0.4954, 0.3420])\n",
      "-- t: 3, cur_and_prev: torch.Size([4, 2]), mean: tensor([ 1.0623, -0.1802])\n",
      "-- t: 4, cur_and_prev: torch.Size([5, 2]), mean: tensor([ 1.1401, -0.4462])\n",
      "-- t: 5, cur_and_prev: torch.Size([6, 2]), mean: tensor([ 1.0870, -0.4071])\n",
      "-- t: 6, cur_and_prev: torch.Size([7, 2]), mean: tensor([ 1.0430, -0.1299])\n",
      "-- t: 7, cur_and_prev: torch.Size([8, 2]), mean: tensor([ 1.1138, -0.1641])\n",
      "tensor([[ 0.1808, -0.0700],\n",
      "        [-0.0894, -0.4926],\n",
      "        [ 0.1490, -0.3199],\n",
      "        [ 0.3504, -0.2238],\n",
      "        [ 0.3525,  0.0545],\n",
      "        [ 0.0688, -0.0396],\n",
      "        [ 0.0927, -0.0682],\n",
      "        [-0.0341,  0.1332]])\n",
      "tensor([[ 0.1808, -0.0700],\n",
      "        [-0.3596, -0.9152],\n",
      "        [ 0.6258,  0.0255],\n",
      "        [ 0.9545,  0.0643],\n",
      "        [ 0.3612,  1.1679],\n",
      "        [-1.3499, -0.5102],\n",
      "        [ 0.2360, -0.2398],\n",
      "        [-0.9211,  1.5433]])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "# A mathematical trick that is used in the self attention inside a transformer\n",
    "# and at the heart of an efficient implementation of self-attention\n",
    "\n",
    "torch.manual_seed(1337)\n",
    "B,T,C = 4,8,2 # batch, time, channels\n",
    "x = torch.randn(B,T,C)\n",
    "print(\"x:\", x.shape)\n",
    "\n",
    "# Now we would like this 8 tokens in a batch to talk to each other. But the token for example at the fifth location\n",
    "# should not communicate with future tokens in a sequence (6, 7, 8, ...). It should only talk to tokens in (4, 3, 2, ...) locations.\n",
    "# So information only flows from previous context to the current timestamp and we cannot get any information from the future\n",
    "# because we are about to try to predict the future.\n",
    "# The easiest way for tokens to communicate is to just do an average of all the preceding elements.\n",
    "# For example if i am the fifth token (T) i would like to take channels (C) that make up information at my step but\n",
    "# then also the channels from the four, third, second and first steps. I'd like to average those up and then that would\n",
    "# become sort of like a feature vector that summarizes me in the context of my history.\n",
    "\n",
    "# for each T inside each B, we wanna calculate the average of current T and all the previous\n",
    "xbow = torch.zeros((B,T,C))\n",
    "for b in range(B):\n",
    "  print(\"---- b:\", b)\n",
    "  for t in range(T):\n",
    "    cur_and_prev = x[b, :t+1]\n",
    "    mean = torch.mean(cur_and_prev, 0)\n",
    "    xbow[b,t] = mean\n",
    "    print(f\"-- t: {t}, cur_and_prev: {cur_and_prev.shape}, mean: {mean}\")\n",
    "\n",
    "print(xbow[0])\n",
    "print(x[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0., 0., 0., 0., 0., 0., 0., 0.],\n",
      "        [0., 0., 0., 0., 0., 0., 0., 0.],\n",
      "        [0., 0., 0., 0., 0., 0., 0., 0.],\n",
      "        [0., 0., 0., 0., 0., 0., 0., 0.],\n",
      "        [0., 0., 0., 0., 0., 0., 0., 0.],\n",
      "        [0., 0., 0., 0., 0., 0., 0., 0.],\n",
      "        [0., 0., 0., 0., 0., 0., 0., 0.],\n",
      "        [0., 0., 0., 0., 0., 0., 0., 0.]])\n",
      "tensor([[0., -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "        [0., 0., -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "        [0., 0., 0., -inf, -inf, -inf, -inf, -inf],\n",
      "        [0., 0., 0., 0., -inf, -inf, -inf, -inf],\n",
      "        [0., 0., 0., 0., 0., -inf, -inf, -inf],\n",
      "        [0., 0., 0., 0., 0., 0., -inf, -inf],\n",
      "        [0., 0., 0., 0., 0., 0., 0., -inf],\n",
      "        [0., 0., 0., 0., 0., 0., 0., 0.]])\n",
      "tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.5000, 0.5000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.3333, 0.3333, 0.3333, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.2500, 0.2500, 0.2500, 0.2500, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.2000, 0.2000, 0.2000, 0.2000, 0.2000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.0000, 0.0000],\n",
      "        [0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.0000],\n",
      "        [0.1250, 0.1250, 0.1250, 0.1250, 0.1250, 0.1250, 0.1250, 0.1250]])\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "# version 2 (the same result) using tril, softmax and matrix mul.\n",
    "\n",
    "tril = torch.tril(torch.ones(T,T))\n",
    "wei = torch.zeros((T,T))\n",
    "print(wei)\n",
    "wei = wei.masked_fill(tril == 0, float(\"-inf\"))\n",
    "print(wei)\n",
    "wei = F.softmax(wei, dim=-1) # normalization to 1\n",
    "print(wei)\n",
    "xbow2 = wei @ x\n",
    "print(torch.allclose(xbow, xbow2))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-- outputs of the dot product is raw affinities between all tokens:\n",
      "tensor([[-1.7629, -1.3011,  0.5652,  2.1616, -1.0674,  1.9632,  1.0765, -0.4530],\n",
      "        [-3.3334, -1.6556,  0.1040,  3.3782, -2.1825,  1.0415, -0.0557,  0.2927],\n",
      "        [-1.0226, -1.2606,  0.0762, -0.3813, -0.9843, -1.4303,  0.0749, -0.9547],\n",
      "        [ 0.7836, -0.8014, -0.3368, -0.8496, -0.5602, -1.1701, -1.2927, -1.0260],\n",
      "        [-1.2566,  0.0187, -0.7880, -1.3204,  2.0363,  0.8638,  0.3719,  0.9258],\n",
      "        [-0.3126,  2.4152, -0.1106, -0.9931,  3.3449, -2.5229,  1.4187,  1.2196],\n",
      "        [ 1.0876,  1.9652, -0.2621, -0.3158,  0.6091,  1.2616, -0.5484,  0.8048],\n",
      "        [-1.8044, -0.4126, -0.8306,  0.5898, -0.7987, -0.5856,  0.6433,  0.6303]],\n",
      "       grad_fn=<SelectBackward0>) \n",
      "\n",
      "-- but we want to ban some tokens from communication (explanation why see at the top of the notebook)\n",
      "tensor([[-1.7629,    -inf,    -inf,    -inf,    -inf,    -inf,    -inf,    -inf],\n",
      "        [-3.3334, -1.6556,    -inf,    -inf,    -inf,    -inf,    -inf,    -inf],\n",
      "        [-1.0226, -1.2606,  0.0762,    -inf,    -inf,    -inf,    -inf,    -inf],\n",
      "        [ 0.7836, -0.8014, -0.3368, -0.8496,    -inf,    -inf,    -inf,    -inf],\n",
      "        [-1.2566,  0.0187, -0.7880, -1.3204,  2.0363,    -inf,    -inf,    -inf],\n",
      "        [-0.3126,  2.4152, -0.1106, -0.9931,  3.3449, -2.5229,    -inf,    -inf],\n",
      "        [ 1.0876,  1.9652, -0.2621, -0.3158,  0.6091,  1.2616, -0.5484,    -inf],\n",
      "        [-1.8044, -0.4126, -0.8306,  0.5898, -0.7987, -0.5856,  0.6433,  0.6303]],\n",
      "       grad_fn=<SelectBackward0>) \n",
      "\n",
      "-- now we use softmax (exponentiate and normalized) because we want to have a nice distribution that sums to one\n",
      "tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.1574, 0.8426, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.2088, 0.1646, 0.6266, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.5792, 0.1187, 0.1889, 0.1131, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.0294, 0.1052, 0.0469, 0.0276, 0.7909, 0.0000, 0.0000, 0.0000],\n",
      "        [0.0176, 0.2689, 0.0215, 0.0089, 0.6812, 0.0019, 0.0000, 0.0000],\n",
      "        [0.1691, 0.4066, 0.0438, 0.0416, 0.1048, 0.2012, 0.0329, 0.0000],\n",
      "        [0.0210, 0.0843, 0.0555, 0.2297, 0.0573, 0.0709, 0.2423, 0.2391]],\n",
      "       grad_fn=<SelectBackward0>)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 8, 16])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# version 3: self-attention!\n",
    "\n",
    "# Different tokens will find different other tokens more or less interesting so each token wants to gather\n",
    "# information from the past in a data depended way, and this is the problem that self-attention solves.\n",
    "# The way self-attention solves it:\n",
    "# Every single token at each position will emit two vectors: query and key.\n",
    "# The query vector is \"what am i looking for?\"\n",
    "# The key vector is \"what do i contain?\"\n",
    "# So the way we get affinities (родство) between these tokens in a sequence: we do a dot product between the keys and the queries.\n",
    "# So my query dot products with all the keys of all the other tokens and that dot product becomes \"wei\" (variable in example below).\n",
    "# So if the key and the query are sort of aligned they will interact to a very high amount and then i will get to learn more about\n",
    "# that specific token as opposed to any other token in a sequence.\n",
    "\n",
    "torch.manual_seed(1337)\n",
    "B,T,C = 4,8,32 # batch, time, channels\n",
    "x = torch.randn(B,T,C)\n",
    "\n",
    "# let's see a single Head perform self-attention\n",
    "head_size = 16\n",
    "\n",
    "# these layers just going to apply a matrix mul with some fixes waights (because bias=False)\n",
    "key = torch.nn.Linear(C, head_size, bias=False)\n",
    "query = torch.nn.Linear(C, head_size, bias=False)\n",
    "value = torch.nn.Linear(C, head_size, bias=False)\n",
    "\n",
    "k = key(x) # [B, T, 16]\n",
    "q = query(x) # [B, T, 16]\n",
    "\n",
    "wei = q @ k.transpose(-2, -1) # [B, T, 16] @ [B, 16, T] ---> [B, T, T]\n",
    "print(\"-- outputs of the dot product is raw affinities between all tokens:\")\n",
    "print(wei[0], \"\\n\")\n",
    "\n",
    "tril = torch.tril(torch.ones(T,T))\n",
    "wei = wei.masked_fill(tril == 0, float(\"-inf\")) # just delete this line if you want to allow all tokens talk to each other fully (ex sentimate analysis)\n",
    "print(\"-- but we want to ban some tokens from communication (explanation why see at the top of the notebook)\")\n",
    "print(wei[0], \"\\n\")\n",
    "wei = F.softmax(wei, dim=-1) # normalization\n",
    "print(\"-- now we use softmax (exponentiate and normalized) because we want to have a nice distribution that sums to one\")\n",
    "print(wei[0])\n",
    "\n",
    "v = value(x)\n",
    "out = wei @ v\n",
    "out.shape\n",
    "\n",
    "# Attention is a communication mechanism. Can be seen as nodes in a directed graph looking at each other and aggregating\n",
    "# information with a weighted sum from all nodes that point to them, with data-dependent weights.\n",
    "\n",
    "# There is no notion of space. Attention simply acts over a set of vectors. This is why we need to positionally encode tokens.\n",
    "\n",
    "# Each example across batch dimension is of course processed completely independently and never \"talk\" to each other.\n",
    "\n",
    "# In an \"encoder\" attention block just delete the single line that does masking with tril, allowing all tokens to communicate.\n",
    "# This block here is called a \"decoder\" attention block because it has triangular masking, and is usually used in autoregressive settings, like language modeling.\n",
    "\n",
    "# \"self-attention\" just means that the keys and values are produced from the same source as queries.\n",
    "# In \"cross-attention\", the queries still get produced from `x`, but the keys and values come from some other, external source (e.g. an encoder module)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# \"Scaled\" attention additionaly divides wei by `1/sqrt(head_size)`. This makes it so when input Q,K are unit variance,\n",
    "# `wei` will be unit variance too and Softmax will stay diffuse and not saturate too much. Illustration below.\n",
    "\n",
    "k = torch.randn(B,T,head_size)\n",
    "q = torch.randn(B,T,head_size)\n",
    "wei = q @ k.transpose(-2, -1) * head_size**-0.5"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}