diff --git a/ch03/01_main-chapter-code/ch03.ipynb b/ch03/01_main-chapter-code/ch03.ipynb index 08c0a40..062efd9 100644 --- a/ch03/01_main-chapter-code/ch03.ipynb +++ b/ch03/01_main-chapter-code/ch03.ipynb @@ -1658,11 +1658,18 @@ "\n", " # Compute scaled dot-product attention (aka self-attention) with a causal mask\n", " attn_scores = queries @ keys.transpose(2, 3) # Dot product for each head\n", - " attn_scores.masked_fill_(self.mask.bool()[:num_tokens, :num_tokens].unsqueeze(0).unsqueeze(0), -torch.inf)\n", + " # Original mask truncated to the number of tokens and converted to boolean\n", + " mask_bool = self.mask.bool()[:num_tokens, :num_tokens]\n", + " # Unsqueeze the mask twice to match dimensions\n", + " mask_unsqueezed = mask_bool.unsqueeze(0).unsqueeze(0)\n", + " # Use the unsqueezed mask to fill attention scores\n", + " attn_scores.masked_fill_(mask_unsqueezed, -torch.inf)\n", + " \n", " attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)\n", " attn_weights = self.dropout(attn_weights)\n", "\n", - " context_vec = (attn_weights @ values).transpose(1, 2) # Shape: (b, num_tokens, n_heads, head_dim)\n", + " # Shape: (b, num_tokens, num_heads, head_dim)\n", + " context_vec = (attn_weights @ values).transpose(1, 2) \n", " \n", " # Combine heads, where self.d_out = self.num_heads * self.head_dim\n", " context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)\n", diff --git a/ch03/01_main-chapter-code/multihead-attention.ipynb b/ch03/01_main-chapter-code/multihead-attention.ipynb index d9432ab..fbbdee8 100644 --- a/ch03/01_main-chapter-code/multihead-attention.ipynb +++ b/ch03/01_main-chapter-code/multihead-attention.ipynb @@ -35,6 +35,7 @@ "source": [ "import tiktoken\n", "import torch\n", + "import torch.nn as nn\n", "from torch.utils.data import Dataset, DataLoader\n", "\n", "\n", @@ -86,8 +87,8 @@ "block_size = max_len\n", "\n", "\n", - "token_embedding_layer = torch.nn.Embedding(vocab_size, output_dim)\n", - "pos_embedding_layer = torch.nn.Embedding(vocab_size, output_dim)\n", + "token_embedding_layer = nn.Embedding(vocab_size, output_dim)\n", + "pos_embedding_layer = nn.Embedding(vocab_size, output_dim)\n", "\n", "max_length = 4\n", "dataloader = create_dataloader(raw_text, batch_size=8, max_length=max_length, stride=5)" @@ -152,17 +153,15 @@ "metadata": {}, "outputs": [], "source": [ - "import torch\n", - "\n", - "class CausalSelfAttention(torch.nn.Module):\n", + "class CausalSelfAttention(nn.Module):\n", "\n", " def __init__(self, d_in, d_out, block_size, dropout):\n", " super().__init__()\n", " self.d_out = d_out\n", - " self.W_query = torch.nn.Linear(d_in, d_out, bias=False)\n", - " self.W_key = torch.nn.Linear(d_in, d_out, bias=False)\n", - " self.W_value = torch.nn.Linear(d_in, d_out, bias=False)\n", - " self.dropout = torch.nn.Dropout(dropout) # New\n", + " self.W_query = nn.Linear(d_in, d_out, bias=False)\n", + " self.W_key = nn.Linear(d_in, d_out, bias=False)\n", + " self.W_value = nn.Linear(d_in, d_out, bias=False)\n", + " self.dropout = nn.Dropout(dropout) # New\n", " self.register_buffer('mask', torch.triu(torch.ones(block_size, block_size), diagonal=1)) # New\n", "\n", " def forward(self, x):\n", @@ -181,14 +180,14 @@ " return context_vec\n", "\n", "\n", - "class MultiHeadAttentionWrapper(torch.nn.Module):\n", + "class MultiHeadAttentionWrapper(nn.Module):\n", " def __init__(self, d_in, d_out, block_size, dropout, num_heads):\n", " super().__init__()\n", - " self.heads = torch.nn.ModuleList(\n", + " self.heads = nn.ModuleList(\n", " [CausalSelfAttention(d_in, d_out, block_size, dropout) \n", " for _ in range(num_heads)]\n", " )\n", - " self.out_proj = torch.nn.Linear(d_out*num_heads, d_out*num_heads)\n", + " self.out_proj = nn.Linear(d_out*num_heads, d_out*num_heads)\n", "\n", " def forward(self, x):\n", " context_vec = torch.cat([head(x) for head in self.heads], dim=-1)\n", @@ -241,10 +240,7 @@ "metadata": {}, "outputs": [], "source": [ - "import torch\n", - "\n", - "\n", - "class MultiHeadAttention(torch.nn.Module):\n", + "class MultiHeadAttention(nn.Module):\n", " def __init__(self, d_in, d_out, block_size, dropout, num_heads):\n", " super().__init__()\n", " assert d_out % num_heads == 0, \"d_out must be divisible by n_heads\"\n", @@ -253,30 +249,48 @@ " self.num_heads = num_heads\n", " self.head_dim = d_out // num_heads # Reduce the projection dim to match desired output dim\n", "\n", - " self.W_query = torch.nn.Linear(d_in, d_out, bias=False)\n", - " self.W_key = torch.nn.Linear(d_in, d_out, bias=False)\n", - " self.W_value = torch.nn.Linear(d_in, d_out, bias=False)\n", - " self.out_proj = torch.nn.Linear(d_out, d_out) # Linear layer to combine head outputs\n", - " self.dropout = torch.nn.Dropout(dropout)\n", + " self.W_query = nn.Linear(d_in, d_out, bias=False)\n", + " self.W_key = nn.Linear(d_in, d_out, bias=False)\n", + " self.W_value = nn.Linear(d_in, d_out, bias=False)\n", + " self.out_proj = nn.Linear(d_out, d_out) # Linear layer to combine head outputs\n", + " self.dropout = nn.Dropout(dropout)\n", " self.register_buffer('mask', torch.triu(torch.ones(block_size, block_size), diagonal=1))\n", "\n", " def forward(self, x):\n", - " b, n_tokens, d_in = x.shape\n", + " b, num_tokens, d_in = x.shape\n", "\n", - " # Split into multiple heads\n", - " keys = self.W_key(x).view(b, n_tokens, self.num_heads, self.head_dim).transpose(1, 2)\n", - " queries = self.W_query(x).view(b, n_tokens, self.num_heads, self.head_dim).transpose(1, 2)\n", - " values = self.W_value(x).view(b, n_tokens, self.num_heads, self.head_dim).transpose(1, 2)\n", + " keys = self.W_key(x) # Shape: (b, num_tokens, d_out)\n", + " queries = self.W_query(x)\n", + " values = self.W_value(x)\n", "\n", - " # Compute scaled dot-product attention for each head\n", + " # We implicitly split the matrix by adding a `num_heads` dimension\n", + " # Unroll last dim: (b, num_tokens, d_out) -> (b, num_tokens, num_heads, head_dim)\n", + " keys = keys.view(b, num_tokens, self.num_heads, self.head_dim) \n", + " values = values.view(b, num_tokens, self.num_heads, self.head_dim)\n", + " queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)\n", + "\n", + " # Transpose: (b, num_tokens, num_heads, head_dim) -> (b, num_heads, num_tokens, head_dim)\n", + " keys = keys.transpose(1, 2)\n", + " queries = queries.transpose(1, 2)\n", + " values = values.transpose(1, 2)\n", + "\n", + " # Compute scaled dot-product attention (aka self-attention) with a causal mask\n", " attn_scores = queries @ keys.transpose(2, 3) # Dot product for each head\n", - " attn_scores.masked_fill_(self.mask.bool()[:n_tokens, :n_tokens].unsqueeze(0).unsqueeze(0), -torch.inf)\n", + " # Original mask truncated to the number of tokens and converted to boolean\n", + " mask_bool = self.mask.bool()[:num_tokens, :num_tokens]\n", + " # Unsqueeze the mask twice to match dimensions\n", + " mask_unsqueezed = mask_bool.unsqueeze(0).unsqueeze(0)\n", + " # Use the unsqueezed mask to fill attention scores\n", + " attn_scores.masked_fill_(mask_unsqueezed, -torch.inf)\n", + " \n", " attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)\n", " attn_weights = self.dropout(attn_weights)\n", - " context_vec = (attn_weights @ values).transpose(1, 2) # Shape: (b, T, n_heads, head_dim)\n", + "\n", + " # Shape: (b, num_tokens, num_heads, head_dim)\n", + " context_vec = (attn_weights @ values).transpose(1, 2) \n", " \n", " # Combine heads, where self.d_out = self.num_heads * self.head_dim\n", - " context_vec = context_vec.contiguous().view(b, n_tokens, self.d_out)\n", + " context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)\n", " context_vec = self.out_proj(context_vec) # optional projection\n", "\n", " return context_vec"