self-attention.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 公式：\n",
    "$$\n",
    "Attention(Q, K, V) = Softmax(\\frac{QK^T}{\\sqrt{d_k}})V\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 第一重境界：简化版本"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "class SelfAttentionV1(nn.Module):\n",
    "    def __init__(self, hidden_dim: int =728) -> None:\n",
    "        super().__init__()\n",
    "        self.hidden_dim = hidden_dim\n",
    "\n",
    "        self.query_proj = nn.Linear(hidden_dim, hidden_dim)\n",
    "        self.key_proj = nn.Linear(hidden_dim, hidden_dim)    \n",
    "        self.value_proj = nn.Linear(hidden_dim, hidden_dim)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        # x: [batch_size, seq_len, hidden_dim]\n",
    "        Q = self.query_proj(x)\n",
    "        K = self.key_proj(x)\n",
    "        V = self.value_proj(x)\n",
    "        # Q K V: [batch_size, seq_len, hidden_dim]\n",
    "\n",
    "        attention_value = torch.matmul(Q, K.transpose(1, 2)) / math.sqrt(self.hidden_dim)\n",
    "        # attention_value: [batch_size, seq_len, seq_len]\n",
    "        attention_weight = torch.softmax(attention_value, dim=-1)\n",
    "        print(attention_weight)\n",
    "        # attention_weight: [batch_size, seq_len, seq_len]\n",
    "        output = torch.matmul(attention_weight, V)\n",
    "        # output: [batch_size, seq_len, hidden_dim]\n",
    "\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = torch.randn(3, 2, 4)\n",
    "print(X)\n",
    "\n",
    "self_att_net = SelfAttentionV1(4)\n",
    "self_att_net(X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 探究softmax的指定dim效果\n",
    "X = torch.randn(3, 2, 4)\n",
    "print(X)\n",
    "Y = torch.softmax(X, dim=1)\n",
    "print(Y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 第二重境界：效率优化\n",
    "只做一次proj"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "class SelfAttentionV2(nn.Module):\n",
    "    def __init__(self, hidden_dim: int =728) -> None:\n",
    "        super().__init__()\n",
    "        self.hidden_dim = hidden_dim\n",
    "\n",
    "        self.proj = nn.Linear(hidden_dim, hidden_dim * 3)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        # x: [batch_size, seq_len, hidden_dim]\n",
    "        QKV = self.proj(x)\n",
    "        # QKV: [batch_size, seq_len, hidden_dim * 3]\n",
    "        Q, K, V = torch.split(QKV, self.hidden_dim, dim=-1)\n",
    "\n",
    "        attention_value = torch.matmul(Q, K.transpose(1, 2)) / math.sqrt(self.hidden_dim)\n",
    "        # attention_value: [batch_size, seq_len, seq_len]\n",
    "        attention_weight = torch.softmax(attention_value, dim=-1)\n",
    "        # attention_weight: [batch_size, seq_len, seq_len]\n",
    "        output = attention_weight @ V  # @ 等于torch.matmul()\n",
    "        # output: [batch_size, seq_len, hidden_dim]\n",
    "\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = torch.randn(3, 2, 4)\n",
    "print(X)\n",
    "\n",
    "self_att_net = SelfAttentionV2(4)\n",
    "self_att_net(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 第三重境界：加入一些细节\n",
    "* dropout 的位置\n",
    "* attention_mask：输入的seq_len不可能都一样\n",
    "* MHA 中还有一个output_proj映射"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "class SelfAttentionV3(nn.Module):\n",
    "    def __init__(self, hidden_dim: int =728, dropout_rate=0.1) -> None:\n",
    "        super().__init__()\n",
    "        self.hidden_dim = hidden_dim\n",
    "\n",
    "        self.query_proj = nn.Linear(hidden_dim, hidden_dim)\n",
    "        self.key_proj = nn.Linear(hidden_dim, hidden_dim)    \n",
    "        self.value_proj = nn.Linear(hidden_dim, hidden_dim)\n",
    "\n",
    "        self.attention_dropout = nn.Dropout(dropout_rate)\n",
    "\n",
    "        self.output_proj = nn.Linear(hidden_dim, hidden_dim)\n",
    "    \n",
    "    def forward(self, x, attention_mask=None):\n",
    "        # x: [batch_size, seq_len, hidden_dim]\n",
    "        Q = self.query_proj(x)\n",
    "        K = self.key_proj(x)\n",
    "        V = self.value_proj(x)\n",
    "        # Q K V: [batch_size, seq_len, hidden_dim]\n",
    "\n",
    "        attention_value = torch.matmul(Q, K.transpose(1, 2)) / math.sqrt(self.hidden_dim)\n",
    "        # attention_value: [batch_size, seq_len, seq_len]\n",
    "\n",
    "        if attention_mask is not None:\n",
    "            # 根据给定的条件（attention_mask == 0）对attention_value进行修改，将满足条件的元素替换为指定的值\n",
    "            attention_value = attention_value.masked_fill(attention_mask == 0, -1e9)\n",
    "\n",
    "        attention_weight = torch.softmax(attention_value, dim=-1)\n",
    "        print(\"attention_weight:\", attention_weight)\n",
    "\n",
    "        attention_weight = self.attention_dropout(attention_weight)  # dropout在这里！！！将有些词直接drop掉\n",
    "        print(\"attention_weight:\", attention_weight)\n",
    "        # attention_weight: [batch_size, seq_len, seq_len]\n",
    "\n",
    "        attention_result = torch.matmul(attention_weight, V)\n",
    "        # attention_result: [batch_size, seq_len, hidden_dim]\n",
    "\n",
    "        output = self.output_proj(attention_result)\n",
    "\n",
    "        return output "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[1, 1, 1, 0],\n",
      "         [1, 1, 1, 0],\n",
      "         [1, 1, 1, 0],\n",
      "         [1, 1, 1, 0]],\n",
      "\n",
      "        [[1, 1, 0, 0],\n",
      "         [1, 1, 0, 0],\n",
      "         [1, 1, 0, 0],\n",
      "         [1, 1, 0, 0]],\n",
      "\n",
      "        [[1, 0, 0, 0],\n",
      "         [1, 0, 0, 0],\n",
      "         [1, 0, 0, 0],\n",
      "         [1, 0, 0, 0]]])\n"
     ]
    }
   ],
   "source": [
    "x = torch.randn(3, 4, 2)\n",
    "# x : (batch_size, seq_len, hidden_dim)\n",
    "mask = torch.tensor([[1, 1, 1, 0], [1, 1, 0, 0], [1, 0, 0, 0]])\n",
    "# mask : (batch_size, seq_len)\n",
    "# 要和attention_weight的shape保持一致：(batch_size, seq_len, seq_len)\n",
    "# 扩维\n",
    "mask = mask.unsqueeze(dim=1).repeat(1, x.size(1), 1)\n",
    "# unsqueeze在指定的维度（dim=1）上增加一个新的维度，变成(batch_size, 1, seq_len)\n",
    "# repeat沿着指定的维度重复张量，变成(batch_size, seq_len, seq_len)\n",
    "print(mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1, 1, 1, 0],\n",
      "        [1, 1, 0, 0],\n",
      "        [1, 0, 0, 0]]) torch.Size([3, 4])\n",
      "tensor([[[1, 1, 1, 0]],\n",
      "\n",
      "        [[1, 1, 0, 0]],\n",
      "\n",
      "        [[1, 0, 0, 0]]]) torch.Size([3, 1, 4])\n",
      "tensor([[[1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1, 0]],\n",
      "\n",
      "        [[1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0]],\n",
      "\n",
      "        [[1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0]]]) torch.Size([3, 1, 16])\n",
      "tensor([[[[1, 1, 1, 0],\n",
      "          [1, 1, 1, 0],\n",
      "          [1, 1, 1, 0],\n",
      "          [1, 1, 1, 0]]],\n",
      "\n",
      "\n",
      "        [[[1, 1, 0, 0],\n",
      "          [1, 1, 0, 0],\n",
      "          [1, 1, 0, 0],\n",
      "          [1, 1, 0, 0]]],\n",
      "\n",
      "\n",
      "        [[[1, 0, 0, 0],\n",
      "          [1, 0, 0, 0],\n",
      "          [1, 0, 0, 0],\n",
      "          [1, 0, 0, 0]]]]) torch.Size([3, 1, 4, 4])\n"
     ]
    }
   ],
   "source": [
    "# 深入理解unsqueeze和repeat\n",
    "mask = torch.tensor([[1, 1, 1, 0], [1, 1, 0, 0], [1, 0, 0, 0]])\n",
    "print(mask, mask.size())\n",
    "mask = mask.unsqueeze(dim=1)\n",
    "print(mask, mask.size())\n",
    "mask = mask.repeat(1, 1, x.size(1))\n",
    "print(mask, mask.size())\n",
    "mask = mask.reshape(x.size(0), 1, x.size(1), x.size(1))\n",
    "print(mask, mask.size())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "attention_weight: tensor([[[[0.4364, 0.4057, 0.1580, 0.0000],\n",
      "          [0.3586, 0.3635, 0.2778, 0.0000],\n",
      "          [0.3707, 0.3408, 0.2885, 0.0000],\n",
      "          [0.3331, 0.3058, 0.3611, 0.0000]],\n",
      "\n",
      "         [[0.3782, 0.4063, 0.2155, 0.0000],\n",
      "          [0.2360, 0.2727, 0.4914, 0.0000],\n",
      "          [0.2831, 0.2639, 0.4530, 0.0000],\n",
      "          [0.4154, 0.4344, 0.1501, 0.0000]],\n",
      "\n",
      "         [[0.2569, 0.4696, 0.2735, 0.0000],\n",
      "          [0.1911, 0.5683, 0.2407, 0.0000],\n",
      "          [0.3203, 0.3579, 0.3218, 0.0000],\n",
      "          [0.2732, 0.4304, 0.2963, 0.0000]]],\n",
      "\n",
      "\n",
      "        [[[0.5182, 0.4818, 0.0000, 0.0000],\n",
      "          [0.4966, 0.5034, 0.0000, 0.0000],\n",
      "          [0.5211, 0.4789, 0.0000, 0.0000],\n",
      "          [0.5214, 0.4786, 0.0000, 0.0000]],\n",
      "\n",
      "         [[0.4821, 0.5179, 0.0000, 0.0000],\n",
      "          [0.4639, 0.5361, 0.0000, 0.0000],\n",
      "          [0.5176, 0.4824, 0.0000, 0.0000],\n",
      "          [0.4888, 0.5112, 0.0000, 0.0000]],\n",
      "\n",
      "         [[0.3536, 0.6464, 0.0000, 0.0000],\n",
      "          [0.2516, 0.7484, 0.0000, 0.0000],\n",
      "          [0.4723, 0.5277, 0.0000, 0.0000],\n",
      "          [0.3883, 0.6117, 0.0000, 0.0000]]],\n",
      "\n",
      "\n",
      "        [[[1.0000, 0.0000, 0.0000, 0.0000],\n",
      "          [1.0000, 0.0000, 0.0000, 0.0000],\n",
      "          [1.0000, 0.0000, 0.0000, 0.0000],\n",
      "          [1.0000, 0.0000, 0.0000, 0.0000]],\n",
      "\n",
      "         [[1.0000, 0.0000, 0.0000, 0.0000],\n",
      "          [1.0000, 0.0000, 0.0000, 0.0000],\n",
      "          [1.0000, 0.0000, 0.0000, 0.0000],\n",
      "          [1.0000, 0.0000, 0.0000, 0.0000]],\n",
      "\n",
      "         [[1.0000, 0.0000, 0.0000, 0.0000],\n",
      "          [1.0000, 0.0000, 0.0000, 0.0000],\n",
      "          [1.0000, 0.0000, 0.0000, 0.0000],\n",
      "          [1.0000, 0.0000, 0.0000, 0.0000]]]], grad_fn=<SoftmaxBackward0>)\n",
      "attention_weight: tensor([[[[0.4849, 0.4507, 0.1755, 0.0000],\n",
      "          [0.0000, 0.4039, 0.3087, 0.0000],\n",
      "          [0.4119, 0.3786, 0.3206, 0.0000],\n",
      "          [0.3702, 0.3397, 0.0000, 0.0000]],\n",
      "\n",
      "         [[0.4202, 0.4515, 0.2394, 0.0000],\n",
      "          [0.0000, 0.3030, 0.0000, 0.0000],\n",
      "          [0.0000, 0.0000, 0.5033, 0.0000],\n",
      "          [0.4616, 0.4827, 0.1668, 0.0000]],\n",
      "\n",
      "         [[0.2854, 0.5218, 0.3039, 0.0000],\n",
      "          [0.0000, 0.6314, 0.0000, 0.0000],\n",
      "          [0.3559, 0.3977, 0.3575, 0.0000],\n",
      "          [0.3036, 0.4783, 0.3292, 0.0000]]],\n",
      "\n",
      "\n",
      "        [[[0.5758, 0.5353, 0.0000, 0.0000],\n",
      "          [0.5518, 0.5593, 0.0000, 0.0000],\n",
      "          [0.5790, 0.5322, 0.0000, 0.0000],\n",
      "          [0.5794, 0.5318, 0.0000, 0.0000]],\n",
      "\n",
      "         [[0.5356, 0.5755, 0.0000, 0.0000],\n",
      "          [0.5155, 0.5957, 0.0000, 0.0000],\n",
      "          [0.0000, 0.5360, 0.0000, 0.0000],\n",
      "          [0.5431, 0.5680, 0.0000, 0.0000]],\n",
      "\n",
      "         [[0.3929, 0.7182, 0.0000, 0.0000],\n",
      "          [0.2796, 0.8315, 0.0000, 0.0000],\n",
      "          [0.5248, 0.5864, 0.0000, 0.0000],\n",
      "          [0.4315, 0.6797, 0.0000, 0.0000]]],\n",
      "\n",
      "\n",
      "        [[[1.1111, 0.0000, 0.0000, 0.0000],\n",
      "          [1.1111, 0.0000, 0.0000, 0.0000],\n",
      "          [1.1111, 0.0000, 0.0000, 0.0000],\n",
      "          [1.1111, 0.0000, 0.0000, 0.0000]],\n",
      "\n",
      "         [[0.0000, 0.0000, 0.0000, 0.0000],\n",
      "          [1.1111, 0.0000, 0.0000, 0.0000],\n",
      "          [0.0000, 0.0000, 0.0000, 0.0000],\n",
      "          [1.1111, 0.0000, 0.0000, 0.0000]],\n",
      "\n",
      "         [[1.1111, 0.0000, 0.0000, 0.0000],\n",
      "          [1.1111, 0.0000, 0.0000, 0.0000],\n",
      "          [1.1111, 0.0000, 0.0000, 0.0000],\n",
      "          [1.1111, 0.0000, 0.0000, 0.0000]]]], grad_fn=<MulBackward0>)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[[-0.0549,  1.2250],\n",
       "          [-0.1789,  1.1136],\n",
       "          [-0.0378,  1.2962],\n",
       "          [-0.2227,  0.9227]],\n",
       "\n",
       "         [[-0.0359,  1.3616],\n",
       "          [-0.3655,  0.7591],\n",
       "          [-0.2266,  1.1342],\n",
       "          [-0.0459,  1.3197]],\n",
       "\n",
       "         [[-0.0496,  1.2516],\n",
       "          [-0.2866,  0.7378],\n",
       "          [-0.0310,  1.3288],\n",
       "          [-0.0432,  1.2785]]],\n",
       "\n",
       "\n",
       "        [[[-0.0756,  1.1383],\n",
       "          [-0.0749,  1.1425],\n",
       "          [-0.0757,  1.1377],\n",
       "          [-0.0757,  1.1377]],\n",
       "\n",
       "         [[-0.0685,  1.2278],\n",
       "          [-0.0678,  1.2319],\n",
       "          [-0.2751,  0.9257],\n",
       "          [-0.0687,  1.2263]],\n",
       "\n",
       "         [[-0.0759,  1.1355],\n",
       "          [-0.0937,  1.0636],\n",
       "          [-0.0553,  1.2192],\n",
       "          [-0.0699,  1.1600]]],\n",
       "\n",
       "\n",
       "        [[[-0.0916,  1.0449],\n",
       "          [-0.0916,  1.0449],\n",
       "          [-0.0916,  1.0449],\n",
       "          [-0.0916,  1.0449]],\n",
       "\n",
       "         [[-0.4829,  0.5425],\n",
       "          [-0.0858,  1.1106],\n",
       "          [-0.4829,  0.5425],\n",
       "          [-0.0858,  1.1106]],\n",
       "\n",
       "         [[ 0.0366,  1.5912],\n",
       "          [ 0.0366,  1.5912],\n",
       "          [ 0.0366,  1.5912],\n",
       "          [ 0.0366,  1.5912]]]], grad_fn=<ViewBackward0>)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net = SelfAttentionV3(2)\n",
    "net(x, mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 第四重境界：面试写法\n",
    "没啥区别啊，up主又在瞎讲"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "class SelfAttentionV4(nn.Module):\n",
    "    def __init__(self, hidden_dim: int=728, dropout_rate: float=0.1) -> None:\n",
    "        super().__init__()\n",
    "        self.hidden_dim = hidden_dim\n",
    "\n",
    "        self.query_proj = nn.Linear(hidden_dim, hidden_dim)\n",
    "        self.key_proj = nn.Linear(hidden_dim, hidden_dim)    \n",
    "        self.value_proj = nn.Linear(hidden_dim, hidden_dim)\n",
    "\n",
    "        self.attention_dropout = nn.Dropout(dropout_rate)\n",
    "\n",
    "    \n",
    "    def forward(self, x, attention_mask=None):\n",
    "        # x: [batch_size, seq_len, hidden_dim]\n",
    "        Q = self.query_proj(x)\n",
    "        K = self.key_proj(x)\n",
    "        V = self.value_proj(x)\n",
    "        # Q K V: [batch_size, seq_len, hidden_dim]\n",
    "\n",
    "        attention_value = torch.matmul(Q, K.transpose(1, 2)) / math.sqrt(self.hidden_dim)\n",
    "        # attention_value: [batch_size, seq_len, seq_len]\n",
    "\n",
    "        if attention_mask is not None:\n",
    "            # 根据给定的条件（attention_mask == 0）对attention_value进行修改，将满足条件的元素替换为指定的值\n",
    "            attention_value = attention_value.masked_fill(attention_mask == 0, -1e9)\n",
    "\n",
    "        attention_weight = torch.softmax(attention_value, dim=-1)\n",
    "        attention_weight = self.attention_dropout(attention_weight)  # dropout在这里！！！将有些词直接drop掉\n",
    "        # attention_weight: [batch_size, seq_len, seq_len]\n",
    "\n",
    "        output = torch.matmul(attention_weight, V)\n",
    "\n",
    "        return output "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tanshuai",
   "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.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}