Transformer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "def softmax(Z):\n",
    "    Z = np.exp(Z - Z.max(axis=-1, keepdims=True))\n",
    "    return Z / Z.sum(axis=-1, keepdims=True)\n",
    "    \n",
    "def self_attention(X, mask, W_KQV, W_out):\n",
    "    K,Q,V = np.split(X@W_KQV, 3, axis=-1)\n",
    "    attn = softmax(K@Q.swapaxes(-1,-2) / np.sqrt(X.shape[-1]) + mask)\n",
    "    return attn@V@W_out, attn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0., -inf, -inf, -inf, -inf],\n",
       "        [0., 0., -inf, -inf, -inf],\n",
       "        [0., 0., 0., -inf, -inf],\n",
       "        [0., 0., 0., 0., -inf],\n",
       "        [0., 0., 0., 0., 0.]])"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "T = 5\n",
    "M = torch.triu(-float(\"inf\")*torch.ones(T,T),1)\n",
    "M"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "T, d = 100, 64\n",
    "attn = nn.MultiheadAttention(d, 1, bias=False, batch_first=True)\n",
    "# TODO: attn mask, but why this?\n",
    "M = torch.triu(-float(\"inf\") * torch.ones(T, T), 1)\n",
    "X = torch.randn(1, T, d)\n",
    "Y_, A_ = attn(X, X, X, attn_mask=M)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "Y, A = self_attention(X[0].numpy(), M.numpy(), \n",
    "                      attn.in_proj_weight.detach().numpy().T, \n",
    "                      attn.out_proj.weight.detach().numpy().T) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.8144881e-07\n",
      "1.4463458e-06\n"
     ]
    }
   ],
   "source": [
    "print(np.linalg.norm(A - A_[0].detach().numpy()))\n",
    "print(np.linalg.norm(Y - Y_[0].detach().numpy()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Minibatching with BMM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(10, 3, 5, 3)"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# illustration of batch matmul\n",
    "B = np.random.randn(10,3,5,4)\n",
    "C = np.random.randn(10,3,4,3)\n",
    "(B@C).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 10\n",
    "M = torch.triu(-float(\"inf\") * torch.ones(T, T), 1)\n",
    "X = torch.randn(N, T, d)\n",
    "Y_, A_ = attn(X, X, X, attn_mask=M)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "Y, A = self_attention(X.numpy(), M.numpy(), \n",
    "                      attn.in_proj_weight.detach().numpy().T, \n",
    "                      attn.out_proj.weight.detach().numpy().T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5.503326e-07\n",
      "4.5866805e-06\n"
     ]
    }
   ],
   "source": [
    "print(np.linalg.norm(A - A_.detach().numpy()))\n",
    "print(np.linalg.norm(Y - Y_.detach().numpy()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multihead Attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "def multihead_attention(X, mask, heads, W_KQV, W_out):\n",
    "    N, T, d = X.shape\n",
    "    K, Q, V = np.split(X@W_KQV, 3, axis=-1)\n",
    "    K, Q, V = [a.reshape(N, T, heads, d//heads).swapaxes(1, 2) for a in (K, Q, V)]\n",
    "\n",
    "    attn = softmax(K@Q.swapaxes(-1, -2) / np.sqrt(d//heads) + mask)\n",
    "    return (attn@V).swapaxes(1, 2).reshape(N, T, d) @ W_out, attn "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "heads = 4\n",
    "attn = nn.MultiheadAttention(d, heads, bias=False, batch_first=True)\n",
    "Y_, A_ = attn(X, X, X, attn_mask=M)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "Y, A = multihead_attention(X.numpy(), M.numpy(), 4, \n",
    "                           attn.in_proj_weight.detach().numpy().T,\n",
    "                           attn.out_proj.weight.detach().numpy().T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([10, 100, 100])"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A_.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(10, 4, 100, 100)"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4.314429e-06\n",
      "3.8336384e-07\n"
     ]
    }
   ],
   "source": [
    "print(np.linalg.norm(Y - Y_.detach().numpy()))\n",
    "print(np.linalg.norm(A.mean(1) - A_.detach().numpy()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transformer block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "def layer_norm(Z, eps):\n",
    "    return (Z - Z.mean(axis=-1, keepdims=True)) / np.sqrt(Z.var(axis=-1, keepdims=True) + eps)\n",
    "\n",
    "def relu(Z):\n",
    "    return np.maximum(Z, 0)\n",
    "\n",
    "def transformer(X, mask, heads, W_KQV, W_out, W_ff1, W_ff2, eps):\n",
    "    Z = layer_norm(multihead_attention(X, mask, heads, W_KQV, W_out)[0] + X, eps)\n",
    "    return layer_norm(Z + relu(Z@W_ff1)@W_ff2, eps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "trans = nn.TransformerEncoderLayer(d, heads, dim_feedforward=128, dropout=0.0, batch_first=True)\n",
    "trans.linear1.bias.data.zero_()\n",
    "trans.linear2.bias.data.zero_()\n",
    "Y_ = trans(X, M)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "Y = transformer(X.numpy(), M.numpy(), heads, \n",
    "                trans.self_attn.in_proj_weight.detach().numpy().T,\n",
    "                trans.self_attn.out_proj.weight.detach().numpy().T,\n",
    "                trans.linear1.weight.detach().numpy().T,\n",
    "                trans.linear2.weight.detach().numpy().T,\n",
    "                trans.norm1.eps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.807236e-05\n"
     ]
    }
   ],
   "source": [
    "print(np.linalg.norm(Y - Y_.detach().numpy()))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dlsys",
   "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.18"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}