{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "VRdIio8gyD5F" }, "source": [ "# **Assignment 2: Building a Simple Language Model**\n", "\n", "In this assignment, you will build and extend a tiny neural language model step by step. Starting from a single-vector bag-of-words baseline, you will gradually add positional information, non-linear transformations, a simplified form of attention, MLP layers, residual connections, and gating mechanisms to see how each architectural choice changes the model's ability to predict the next token.\n", "\n", "From this assignment, you will know how to construct the core modules of a language model, train it on data, and understand how architectural design choices affect model performance.\n", "\n", "**Key Components You Will Implement**:\n", "\n", "**Cross-Entropy Loss**: Measures next-token prediction accuracy.\\\n", "**Training Loop**: Forward pass, loss computation, and gradient update.\\\n", "**Positional Embeddings**: Encodes token order to enable sequence modeling.\\\n", "**Simplified Attention**: Replaces the equal-weight average with a similarity-weighted one.\\\n", "**Feed-Forward Network (MLP)**: Learning a function of previous token that is more useful than a simple average\\\n", "**Residual Connections**: Stabilizes training and enhances gradient flow.\\\n", "**Gating Mechanism**: Regulates information flow across layers.\n", "\n", "You will write code inside the `TODO` blocks and answer short conceptual questions.\n", "\n", "At the end you will write a paragraph or two discussing your findings\n", "\n", "We will be using a logging framework called TensorBoard to record metrics during training in order to visualize and compare metrics across different architectures.\n", "\n", "A GPU is recommended but not required. You can select the T4 GPU through Colab by going to Runtime > Change Runtime Type > T4 GPU" ] }, { "cell_type": "markdown", "metadata": { "id": "MyQeJ7iYyHIp" }, "source": [ "# **Imports and Dependencies**" ] }, { "cell_type": "code", "source": [ "!pip install pygtrie datasets" ], "metadata": { "id": "3mzYdObhpEgO" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "dK7ek49OuJOv" }, "outputs": [], "source": [ "import itertools\n", "import json\n", "from collections import Counter\n", "from typing import Dict, List\n", "import datasets\n", "import numpy as np\n", "import pygtrie\n", "import torch\n", "import torch.nn as nn\n", "import torch.nn.functional as F\n", "import torch.optim as optim\n", "from torch.utils.data import Dataset, DataLoader, IterableDataset\n", "from torch.utils.tensorboard import SummaryWriter" ] }, { "cell_type": "markdown", "metadata": { "id": "c-8yYry-u3kf" }, "source": [ "# **Tokenizer**\n", "\n", "Here is the same byte-pairing tokenizer that you implemented in A0, except it has been trained with a vocab size of 65536 on linux text data and a special token has been added to identify the end of a document. Try it on a couple example sentences.\n" ] }, { "cell_type": "code", "source": [ "class Tokenizer:\n", " def __init__(self, vocab_path=None, vocab=None, max_token_length=30, partial_trie=None, special_tokens = None):\n", " \"\"\"\n", " special_tokens: optional dict like {'eos_token': '', 'pad_token': ''}\n", " These will be assigned fresh IDs at the end of the vocab and stored separately.\n", " \"\"\"\n", " print(f\"\\nInitializing Tokenizer:\")\n", " print(f\" - vocab_path: {vocab_path}\")\n", " print(f\" - vocab size: {len(vocab) if vocab else 'None'}\")\n", " print(f\" - max_token_length: {max_token_length}\")\n", " self.vocab_path = vocab_path\n", " self.id_to_tok = []\n", " self.trie = pygtrie.Trie()\n", " self.trie = self.trie if partial_trie is None else partial_trie\n", " self.max_token_length = max_token_length\n", "\n", " #adding a special eos token to seperate between rows\n", " self.special_tokens = {}\n", " self.special_id_to_name = {}\n", "\n", " if vocab_path or vocab:\n", " self.update_trie(vocab)\n", "\n", " if special_tokens:\n", " self.add_special_tokens(special_tokens)\n", " @property\n", " def vocab_size(self):\n", " return len(self.id_to_tok) + len(self.special_tokens)\n", "\n", " @property\n", " def eos_token(self):\n", " return self.special_tokens.get('eos_token', (None,None))[0]\n", "\n", " @property\n", " def eos_token_id(self):\n", " return self.special_tokens.get('eos_token', (None,None))[1]\n", "\n", " def add_special_tokens(self, mapping):\n", " \"\"\"\n", " mapping: {'eos_token': '', 'pad_token': '', ...}\n", " Assigns new IDs after the base vocab (not added to trie).\n", " \"\"\"\n", " base = len(self.id_to_tok)\n", " for i, (name, text_token) in enumerate(mapping.items()):\n", " tok_id = base + len(self.special_tokens)\n", " self.special_tokens[name] = (text_token, tok_id)\n", " self.special_id_to_name[tok_id] = name\n", " print(f\"Added special tokens: { {k:v[0] for k,v in self.special_tokens.items()} }\")\n", "\n", " def update_trie(self, new_vocab=None):\n", " print(\"\\nUpdating trie...\")\n", " if new_vocab is None and self.vocab_path:\n", " print(f\"Loading from vocab file: {self.vocab_path}\")\n", " with open(self.vocab_path, 'r') as f:\n", " for i, line in enumerate(f):\n", " token = tuple(json.loads(line)[0])\n", " self.id_to_tok.append(token)\n", " self.trie[token] = i\n", " elif new_vocab:\n", " for token in new_vocab:\n", " print(token)\n", " self.id_to_tok.append(token)\n", " self.trie[token] = len(self.trie)\n", "\n", " def encode(self, text):\n", " return self._tokenize(text, return_ids=True)\n", "\n", " def decode(self, token_ids):\n", " out_bytes = bytearray()\n", " for tid in token_ids:\n", " if 0 <= tid < len(self.id_to_tok):\n", " out_bytes.extend(self.id_to_tok[tid])\n", " # ids outside the range are simply ignored\n", " return out_bytes.decode('utf-8', errors='ignore')\n", "\n", "\n", " def _tokenize(self, text, return_ids=False):\n", " if isinstance(text, str):\n", " text = text.encode('utf-8', errors='ignore')\n", " tokens = []\n", " while text:\n", " prefix = text[:self.max_token_length]\n", " longest = self.trie.longest_prefix(prefix)\n", " if not longest:\n", " print(f\"Warning: No token found for prefix: {prefix[:10]}...\")\n", " break\n", " tokens.append(longest.value if return_ids else longest.key)\n", " text = text[len(longest.key):]\n", " return tokens" ], "metadata": { "id": "tvUeLwUg5lBh" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "import datasets\n", "dataset = datasets.load_dataset('coms4705-hewitt/linuxlike-tokenized', 'default', streaming=True)['train']\n", "with open('vocab-65k-fw-byte.txt', 'w') as f:\n", " for line in dataset:\n", " f.write(line['text'] + '\\n')\n", "tok = Tokenizer(vocab_path='vocab-65k-fw-byte.txt', special_tokens={'eos_token': ''})\n", "print(tok.eos_token, tok.eos_token_id)\n", "print(tok.vocab_size)" ], "metadata": { "id": "dwVgkLGcwG3-" }, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "id": "BCqhVlJQzFPZ" }, "source": [ "# **Single Vector Model**\n", "\n", "In class you've learned some rudimentary ways in which we can build representations of text. Now, we can use it for next token prediction. Here is a very simplified neural language model, on top of which you will be building your own model. It is equivalent to the model presented in Lecture 1: Text Representation and Language Modeling." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "VYnAIRaYzJQy" }, "outputs": [], "source": [ "class SingleVectorWordModel(nn.Module):\n", " def __init__(self, vocab_size, embedding_dim):\n", " \"\"\"\n", " Initialize the model with a single vector per word.\n", "\n", " Args:\n", " vocab_size: Size of the vocabulary\n", " embedding_dim: Dimension of word embeddings\n", " \"\"\"\n", " super().__init__()\n", "\n", " self.embeddings1 = nn.Embedding(vocab_size, embedding_dim)\n", "\n", " # Initialize weights to small values\n", " nn.init.uniform_(self.embeddings1.weight, -0.001, 0.001)\n", "\n", " def encode_context(self, context_indices):\n", " emb1 = self.embeddings1(context_indices) # (bs, seq, d)\n", " return torch.sum(emb1, dim=1) / context_indices.size(1) # (bs, d)\n", "\n", " def score(self, h):\n", " return torch.einsum('vd,bd->bv', self.embeddings1.weight, h)\n", "\n", " def forward(self, context_indices):\n", " \"\"\"\n", " Forward pass of the model.\n", "\n", " Args:\n", " context_indices: Tensor of shape (batch_size, 3) containing indices of context words\n", "\n", " Returns:\n", " logits: Tensor of shape (batch_size, vocab_size) containing prediction logits\n", " \"\"\"\n", " h = self.encode_context(context_indices)\n", " logits = self.score(h)\n", " return logits" ] }, { "cell_type": "markdown", "metadata": { "id": "UTIJSNWEqBcC" }, "source": [ "# **Training**\n", "\n", "Now that we have defined our model, we can start training!\n", "**Tasks**\n", "1. Question: What does the find_similar_words method and print_similar_words method do?\n", "2. Implement the function `cross_entropy_loss` using\n", " `F.log_softmax` and tensor indexing (`torch.gather` or advanced indexing).\n", "3. Fill in the missing steps in `train_model`:\n", " * move tensors to the correct device\n", " * zero gradients\n", " * forward pass → logits\n", " * compute loss\n", " * backward pass\n", " * optimizer & scheduler steps\n", "4. Look at the visualization in Tensorboard and play with the interface.\n", "5. Question: What do you notice about the loss of the first epoch across multiple training runs. Does this surprise you? Why or why not?" ] }, { "cell_type": "code", "source": [ "def find_similar_words(model, tokenizer, query_words: List[str], k: int = 5) -> Dict[str, List[str]]:\n", " embeddings_list = {'emb1': model.embeddings1.weight}\n", " if hasattr(model, 'embeddings2'):\n", " embeddings_list['emb2'] = model.embeddings2.weight\n", " for name, embeddings in embeddings_list.items():\n", " embeddings = F.normalize(embeddings, p=2, dim=1)\n", " results = {}\n", " for query in query_words:\n", " query_id = tokenizer.encode(query)[0]\n", " query_embed = embeddings[query_id]\n", " similarities = torch.matmul(embeddings, query_embed)\n", " top_k = torch.topk(similarities, k=k+1)\n", " similar_words = [tokenizer.decode([idx.item()]) for idx in top_k.indices[1:]]\n", " results[query] = similar_words\n", " yield results\n", "\n", "def print_similar_words(model, tokenizer, query_words: List[str], k: int = 5):\n", " itr = find_similar_words(model, tokenizer, query_words, k)\n", " for results in itr:\n", " print(\"\\nMost similar words:\")\n", " for query, similar in results.items():\n", " print(f\"\\n'{query}':\")\n", " for word in similar:\n", " print(f\" - {word}\")" ], "metadata": { "id": "KryVEh4MtdMT" }, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "source": [ "1. TODO: What does the find_similiar_words method and print_simliar_words method do?\n" ], "metadata": { "id": "iubGSCkwVM0i" } }, { "cell_type": "markdown", "source": [ "\n", "Answer: response here" ], "metadata": { "id": "tJz-xolhWgBt" } }, { "cell_type": "markdown", "source": [ "### Training Objective\n", "\n", "For training a language model, the objective is for the model to assign high probability to the correct next word among all possible candidates in the vocabulary. In order to measure how far part the predicted distribution is from the ground-truth, we use cross entropy loss.\n", "\n", "Cross entropy loss is defined as the negative log-probability of the correct word.\n", "\n", "Intuitively:\n", "- If the model assigns high probability to the correct word, the loss is small.\n", "- If the model assigns low probability, the loss is large.\n", "\n", "Formally, if the model outputs logits of shape $(B, V)$ where\n", "- $B = $ batch size,\n", "- $V = $ vocabulary size,\n", "- $y_i = $ the index of the correct word for example $i$,\n", "\n", "then the loss is\n", "\n", "$$\n", "\\text{CrossEntropyLoss}(\\text{logits}, \\text{targets})\n", "= -\\frac{1}{B} \\sum_{i=1}^{B}\n", "\\log \\frac{\\exp\\big(\\text{logits}_{i,\\, y_i}\\big)}\n", "{\\sum_{j=1}^{V} \\exp\\big(\\text{logits}_{i,j}\\big)}\n", "$$\n", "where $j$ runs over all vocabulary items." ], "metadata": { "id": "DYxdASo4zNlj" } }, { "cell_type": "code", "source": [ "#2. TODO: Implement cross entropy loss using F.log_softmax and torch.gather\n", "\n", "def cross_entropy_loss(logits, targets):\n", " \"\"\"\n", " Compute cross-entropy loss using log and gather.\n", "\n", " Args:\n", " logits: Tensor of shape (batch_size, vocab_size)\n", " targets: Tensor of shape (batch_size,), containing correct word indices\n", "\n", " Returns:\n", " loss: Scalar tensor, mean cross-entropy loss over the batch\n", " \"\"\"\n", " # TODO: Implement log softmax over the vocabulary dimension\n", " log_probs = FILL_IN\n", "\n", " # TODO: Select the log-probabilities corresponding to the target words\n", " chosen_log_probs = FILL_IN\n", "\n", " # TODO: Take the mean negative log likelihood\n", " loss = FILL_IN\n", "\n", " return loss" ], "metadata": { "id": "B4M72EQmYNl9" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "# 3. TODO: fill in the missing parts of train_model\n", "import math\n", "import os\n", "import time\n", "import torch\n", "import torch.nn.functional as F\n", "import torch.optim as optim\n", "from torch.utils.tensorboard import SummaryWriter\n", "from datetime import datetime\n", "\n", "def evaluate(model, data_loader, max_batches=None):\n", " \"\"\"\n", " Average cross-entropy loss and perplexity on held-out data. (Provided, no work required.)\n", " \"\"\"\n", " was_training = model.training\n", " model.eval()\n", " total_loss, total_examples = 0.0, 0\n", " with torch.no_grad():\n", " for b, (context_indices, target_idx) in enumerate(data_loader):\n", " if max_batches is not None and b >= max_batches:\n", " break\n", " context_indices = context_indices.to(model.device)\n", " target_idx = target_idx.to(model.device)\n", " logits = model(context_indices)\n", " total_loss += F.cross_entropy(logits, target_idx, reduction='sum').item()\n", " total_examples += target_idx.size(0)\n", " if was_training:\n", " model.train()\n", " avg_loss = total_loss / total_examples\n", " return avg_loss, math.exp(avg_loss)\n", "\n", "def train_model(model, train_loader, learning_rate, tokenizer=None,\n", " max_steps=3000, experiment_name=None,\n", " val_loader=None, eval_every=1000, eval_batches=100):\n", " \"\"\"\n", " Train model with simple max_steps termination.\n", " \"\"\"\n", " #Creates unique experiment name if not provided\n", " if experiment_name is None:\n", " timestamp = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n", " experiment_name = f\"experiment_{timestamp}\"\n", "\n", " #Initializes TensorBoard writer\n", " log_dir = f\"runs/{experiment_name}\"\n", " if os.path.exists(log_dir):\n", " n = 2\n", " while os.path.exists(f\"{log_dir}_run{n}\"):\n", " n += 1\n", " log_dir = f\"{log_dir}_run{n}\"\n", " print(f\"Log dir already exists; logging to {log_dir} instead\")\n", " writer = SummaryWriter(log_dir=log_dir)\n", "\n", " optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n", "\n", " scheduler = optim.lr_scheduler.LinearLR(\n", " optimizer, start_factor=1.0, end_factor=0.001, total_iters=max_steps\n", " )\n", "\n", " acc_loss = 0.0\n", " acc_count = 0\n", " global_step = 0\n", " best_loss = float('inf')\n", " val_loss, val_ppl = None, None\n", "\n", " writer.add_hparams({\n", " 'learning_rate': learning_rate,\n", " 'max_steps': max_steps,\n", " 'batch_size': train_loader.batch_size,\n", " 'model_params': sum(p.numel() for p in model.parameters()),\n", " }, {})\n", "\n", " t0 = time.time()\n", " model.train()\n", "\n", " print(f\"Training for {max_steps:,} steps...\")\n", "\n", " train_iter = iter(train_loader)\n", "\n", " while global_step < max_steps:\n", " try:\n", " # TODO:\n", " context_indices, target_idx = FILL_IN\n", " except StopIteration:\n", " train_iter = iter(train_loader)\n", " context_indices, target_idx = FILL_IN\n", "\n", " context_indices = context_indices.to(model.device)\n", " target_idx = target_idx.to(model.device)\n", "\n", " # TODO: forward pass\n", " logits = None\n", "\n", " # TODO: compute loss\n", " loss = None\n", "\n", " # TODO: backward pass\n", " # (zero gradients, backpropagate, optimizer step, scheduler step)\n", "\n", "\n", " global_step += 1\n", " acc_loss += loss.item()\n", " acc_count += 1\n", "\n", " # Log to TensorBoard every step\n", " writer.add_scalar('Loss/Train_Step', loss.item(), global_step)\n", " writer.add_scalar('Learning_Rate', scheduler.get_last_lr()[0], global_step)\n", "\n", " if global_step % 100 == 0:\n", " avg_loss = acc_loss / acc_count\n", " current_lr = scheduler.get_last_lr()[0]\n", " elapsed = time.time() - t0\n", " progress = global_step / max_steps * 100\n", "\n", " print(f\"[Step {global_step:,}/{max_steps:,} ({progress:.1f}%)] Loss: {avg_loss:.4f} | LR: {current_lr:.2e} | Time: {elapsed:.1f}s\")\n", "\n", " writer.add_scalar('Loss/Train_Avg100', avg_loss, global_step)\n", "\n", " if avg_loss < best_loss:\n", " best_loss = avg_loss\n", " writer.add_scalar('Best_Loss', best_loss, global_step)\n", "\n", " acc_loss = 0.0\n", " acc_count = 0\n", "\n", " if tokenizer is not None and global_step % 500 == 0:\n", " try:\n", " query_words = [\"linux\", \"python\", \"code\", \"system\", \"file\"]\n", " print_similar_words(model, tokenizer, query_words, k=3)\n", " except Exception as e:\n", " print(f\"Similarity check failed: {e}\")\n", "\n", " # Evaluate on the validation set occasionally\n", " if val_loader is not None and global_step % eval_every == 0:\n", " val_loss, val_ppl = evaluate(model, val_loader, max_batches=eval_batches)\n", " writer.add_scalar('Loss/Val', val_loss, global_step)\n", " writer.add_scalar('Perplexity/Val', val_ppl, global_step)\n", " print(f\"[Step {global_step:,}] Val loss: {val_loss:.4f} | Val perplexity: {val_ppl:.2f}\")\n", "\n", " # Save checkpoint occasionally\n", " if global_step % 2000 == 0:\n", " checkpoint_path = f\"model_checkpoint_step_{global_step}.pt\"\n", " torch.save({\n", " 'step': global_step,\n", " 'model_state_dict': model.state_dict(),\n", " 'optimizer_state_dict': optimizer.state_dict(),\n", " 'loss': best_loss,\n", " }, checkpoint_path)\n", " print(f\"Saved checkpoint: {checkpoint_path}\")\n", "\n", " if val_loader is not None and global_step % eval_every != 0:\n", " val_loss, val_ppl = evaluate(model, val_loader, max_batches=eval_batches)\n", " writer.add_scalar('Loss/Val', val_loss, global_step)\n", " writer.add_scalar('Perplexity/Val', val_ppl, global_step)\n", "\n", " writer.close()\n", "\n", " final_time = time.time() - t0\n", " print(f\"\\nTraining completed!\")\n", " print(f\"Total steps: {global_step:,}\")\n", " print(f\"Total time: {final_time:.1f}s\")\n", " print(f\"Best loss achieved: {best_loss:.4f}\")\n", " if val_loss is not None:\n", " print(f\"Final val loss: {val_loss:.4f} | Final val perplexity: {val_ppl:.2f}\")\n", " print(f\"TensorBoard logs saved to: {log_dir}\")\n", " print(f\"To view: tensorboard --logdir={log_dir}\")\n", "\n", " return {\n", " 'final_step': global_step,\n", " 'final_loss': best_loss,\n", " 'total_time': final_time,\n", " 'experiment_name': experiment_name,\n", " 'val_loss': val_loss,\n", " 'val_perplexity': val_ppl,\n", " }" ], "metadata": { "id": "AiGLxnACdOgp" }, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "source": [ "Training Setup" ], "metadata": { "id": "5-j-Rna5VpMj" } }, { "cell_type": "markdown", "source": [ "The preprocessing first **tokenizes each document**, optionally truncates it to a maximum length, and appends an **end-of-sequence (EOS)** token to mark document boundaries. \n", "All tokenized documents are **concatenated into a single 1-D array of token IDs**, forming one long text stream.\n", "\n", "The `PackedLMDataset` then **slides a fixed-length context window** of size `T` across this stream. \n", "For each valid position, it outputs:\n", "* **x** – the next `T` tokens (the context)\n", "* **y** – the single token immediately following that context (the prediction target)\n", "\n", "This setup trains the model to predict the next token given a rolling window of prior tokens, even across document boundaries.\n", "\n", "Preprocessing for 100k documents should be sufficient for training, expect it to take around 6 minutes on colab. No work required here.\n" ], "metadata": { "id": "c7sWOXIDI_uS" } }, { "cell_type": "code", "source": [ "from tqdm import tqdm\n", "class PackedLMDataset(Dataset):\n", " \"\"\"\n", " token_ids: 1-D torch.LongTensor of packed tokens (… doc + EOS + doc + EOS …)\n", " context_size: length T of the context window.\n", " Returns (x:[T], y:int)\n", " \"\"\"\n", " def __init__(self, token_ids: torch.Tensor, context_size: int):\n", " assert token_ids.dim() == 1\n", " self.ids = token_ids\n", " self.T = int(context_size)\n", " self.N = max(0, len(self.ids) - self.T) # valid start positions\n", "\n", " def __len__(self):\n", " return self.N\n", "\n", " def __getitem__(self, i):\n", " x = self.ids[i : i + self.T]\n", " y = self.ids[i + self.T]\n", " return x.long(), y.long()\n", "\n", "def pack_rows(rows, tokenizer, max_doc_tokens, eos_id):\n", " \"\"\"\n", " Tokenize each row, truncate to max_doc_tokens, append EOS, and concatenate.\n", " Returns a 1D numpy array of dtype int64.\n", " \"\"\"\n", " pieces = []\n", " append = pieces.append\n", "\n", " for r in tqdm(rows, total=len(rows)):\n", " t = r.get(\"text\") if isinstance(r, dict) else r[\"text\"]\n", " if not t:\n", " continue\n", "\n", " ids = tokenizer.encode(t)\n", " if not ids:\n", " continue\n", "\n", " if max_doc_tokens is not None:\n", " ids = ids[:max_doc_tokens]\n", "\n", " append(np.asarray(ids, dtype=np.int64))\n", " append(np.asarray([eos_id], dtype=np.int64))\n", "\n", " if not pieces:\n", " return np.zeros((1,), dtype=np.int64)\n", "\n", " return np.concatenate(pieces, axis=0)" ], "metadata": { "id": "qg2D-bK62-wZ" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "# Hyperparameters and dataset loading\n", "CONTEXT = 10\n", "MAX_DOC_TOKENS = 2500\n", "VAL_FRACTION = 0.05\n", "EOS_ID = len(tok.id_to_tok)\n", "VOCAB_SIZE = len(tok.id_to_tok) + 1\n", "NUM_DOCS = 100000\n", "\n", "print(\"Loading dataset...\")\n", "ds_full = datasets.load_dataset('coms4705-hewitt/fineweb-linuxlike', 'default', split='train')\n", "ds_subset = ds_full.shuffle(seed=123).select(range(NUM_DOCS))\n", "splits = ds_subset.train_test_split(test_size=VAL_FRACTION, seed=123)\n", "train_rows, val_rows = splits['train'], splits['test']\n", "\n", "print(\"Packing training data...\")\n", "train_ids_np = pack_rows(train_rows, tok, MAX_DOC_TOKENS, EOS_ID)\n", "print(\"Packing validation data...\")\n", "val_ids_np = pack_rows(val_rows, tok, MAX_DOC_TOKENS, EOS_ID)\n", "\n", "train_ids = torch.from_numpy(train_ids_np)\n", "val_ids = torch.from_numpy(val_ids_np)\n", "\n", "train_ds = PackedLMDataset(train_ids, context_size=CONTEXT)\n", "val_ds = PackedLMDataset(val_ids, context_size=CONTEXT)\n", "\n", "train_loader = DataLoader(\n", " train_ds, batch_size=512, shuffle=True,\n", " num_workers=0, drop_last=True\n", ")\n", "val_loader = DataLoader(\n", " val_ds, batch_size=512, shuffle=False,\n", " num_workers=0, drop_last=False\n", ")\n", "\n", "print(f\"Packed tokens: train={len(train_ids):,} val={len(val_ids):,}\")\n", "print(f\"Training examples: {len(train_ds):,}\")\n", "print(f\"Validation examples: {len(val_ds):,}\")\n", "print(f\"Vocab size for the model (include EOS): {VOCAB_SIZE}\")" ], "metadata": { "id": "wAiFcGFBHAIk" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "EMBEDDING_DIM = 100\n", "learning_rate = 0.001" ], "metadata": { "id": "Uo551j3TJZ1Q" }, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "source": [ "# **Time to train your first model!**\n", "\n", "notes: Please do not change the model names [single_vector_model, positional_model, mlpdeep, residualmlp, gatedresidualmlp]" ], "metadata": { "id": "fFsdJ-XAQkve" } }, { "cell_type": "code", "source": [ "#TRAINING THE MODEL (should approximately take 2 minutes on the T4 GPU for 3000 steps)\n", "single_vector_model = SingleVectorWordModel(VOCAB_SIZE, EMBEDDING_DIM)\n", "single_vector_model.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "single_vector_model.to(single_vector_model.device)\n", "\n", "print('Model has {} parameters'.format(sum([np.prod(p.size()) for p in single_vector_model.parameters()])))\n", "\n", "history = train_model(single_vector_model, train_loader, learning_rate, tokenizer=tok, max_steps=3000, experiment_name = \"SingleVectorWordModel1\", val_loader=val_loader)" ], "metadata": { "id": "7Lu_bHuIMilk" }, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "source": [ "4. TODO: You can use Tensorboard like this to visualize your training metrics" ], "metadata": { "id": "qz3QFlOYdiZ4" } }, { "cell_type": "code", "source": [ "%load_ext tensorboard\n", "%tensorboard --logdir runs" ], "metadata": { "id": "JSj-cAW5Mfpj" }, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "source": [ "5. TODO: Try training the model multiple times. What do you notice about the loss of the first step across multiple training runs. Does this surprise you? Why or why not?\n", "\n", "Answer: response here" ], "metadata": { "id": "4vBwwwBxndto" } }, { "cell_type": "markdown", "source": [ "# **Improving the model**\n", "Let's improve the model by allowing it to capture the positional meaning in a text sequence.\n", "$$\n", "h_{< i} =\n", "\\sigma \\left(\n", " \\frac{1}{\\,i-1\\,}\n", " \\sum_{j=1}^{\\,i-1}\n", " \\sigma\\big(E x_j + p_j\\big)\n", "\\right)\n", "$$\n", "\n", "Task: Complete the TODO: sections in the NonLinearPositionalContextModel class" ], "metadata": { "id": "Q7fCeWrj7d-d" } }, { "cell_type": "code", "source": [ "class NonLinearPositionalContextModel(SingleVectorWordModel):\n", " def __init__(self, vocab_size, embedding_dim, max_context):\n", " super().__init__(vocab_size, embedding_dim)\n", " #-------Begin TODO-------\n", " self.max_context = FILL_IN\n", " self.pos = nn.Embedding(FILL_IN, FILL_IN)\n", " nn.init.uniform_(FILL_IN, FILL_IN, FILL_IN)\n", " self.sigmoid = FILL_IN\n", " #-------End TODO-----------\n", "\n", " def encode_context(self, context_indices):\n", " #--------Begin TODO--------\n", " return FILL_IN\n", " #--------End TODO----------\n" ], "metadata": { "id": "N1MbwKNC7cxq" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "#TRAINING THE MODEL\n", "positional_model = NonLinearPositionalContextModel(VOCAB_SIZE, EMBEDDING_DIM, max_context=10)\n", "positional_model.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "positional_model.to(positional_model.device)\n", "\n", "print('Model has {} parameters'.format(sum([np.prod(p.size()) for p in positional_model.parameters()])))\n", "\n", "history = train_model(positional_model, train_loader, learning_rate, tokenizer=tok, max_steps=3000, experiment_name = \"NonLinearPositionalContextModel1\", val_loader=val_loader)" ], "metadata": { "id": "pfnpzVP-pbHx" }, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "source": [ "### From averaging to attention\n", "\n", "The positional model gives every context token the same weight, $1/L$. Attention replaces\n", "that equal-weight average with a weighted one, where each token decides how much of every\n", "other token to pull in — tokens that are similar pull more from each other. See Lecture 3,\n", "\"Information flow and simple attention,\" for the derivation.\n", "\n", "Let $Z \\in \\mathbb{R}^{B \\times L \\times d}$ be the token representations the positional\n", "model already builds, $z_j = \\sigma(Ex_j + p_j)$. The whole operation is three lines:\n", "\n", "$$e = \\frac{ZZ^\\top}{\\sqrt{d}} \\qquad \\alpha = \\text{softmax}(e) \\qquad H = \\alpha Z$$\n", "\n", "with $e, \\alpha \\in \\mathbb{R}^{B \\times L \\times L}$ and $H \\in \\mathbb{R}^{B \\times L \\times d}$.\n", "Remember that $ZZ^\\top$ computes $e_{bij} = z_{bi}^\\top z_{bj}$, and the softmax is taken over the **last** axis —\n", "of the two length-$L$ axes, the first indexes the token doing the attending and the second\n", "the tokens being attended to. The $\\sqrt{d}$ keeps the scores from growing with $d$; more\n", "on that in Lecture 8.\n", "\n", "Task: Complete the TODO: sections in the SimplifiedAttentionContextModel class" ], "metadata": { "id": "JJmgt9ryalsS" } }, { "cell_type": "code", "source": [ "class SimplifiedAttentionContextModel(NonLinearPositionalContextModel):\n", " \"\"\"\n", " Same embeddings, positions, nonlinearity, and scoring as the positional model.\n", " The only change: the equal-weight average becomes a similarity-weighted one.\n", " \"\"\"\n", " def encode_context(self, context_indices):\n", " B, L = context_indices.size()\n", " we = self.embeddings1(context_indices) # (B, L, d)\n", " pe = self.pos(torch.arange(L, device=context_indices.device))[None] # (1, L, d)\n", " Z = self.sigmoid(we + pe) # (B, L, d)\n", " # Note that each score is a sum of d terms. What happens to the softmax if d grows and we don't scale?\n", " scale = Z.size(-1) ** 0.5 # keeps scores from growing with d.\n", "\n", " #--------Begin TODO--------\n", " # 1) all pairwise similarities e[b, i, j] = Z[b, i] . Z[b, j],\n", " # divided by scale -> (B, L, L)\n", " e = FILL_IN\n", " # 2) normalize over the tokens being attended to, the last axis -> (B, L, L)\n", " alpha = FILL_IN\n", " # 3) weighted sum H[b, i] = sum_j alpha[b, i, j] * Z[b, j] -> (B, L, d)\n", " H = FILL_IN\n", " #--------End TODO----------\n", "\n", " self.last_alpha = alpha.detach() # saved so you can inspect the weights below\n", " return self.sigmoid(H[:, -1]) # last position's representation, (B, d)" ], "metadata": { "id": "5L2y8fIPalsS" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "#TRAINING THE MODEL\n", "attention_model = SimplifiedAttentionContextModel(VOCAB_SIZE, EMBEDDING_DIM, max_context=10)\n", "attention_model.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "attention_model.to(attention_model.device)\n", "\n", "print('Model has {} parameters'.format(sum([np.prod(p.size()) for p in attention_model.parameters()])))\n", "\n", "history = train_model(attention_model, train_loader, learning_rate, tokenizer=tok, max_steps=3000, experiment_name = \"SimplifiedAttentionContextModel1\", val_loader=val_loader)" ], "metadata": { "id": "qbuWm6J1alsS" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "# The attention weights are saved in model.last_alpha during the forward pass.\n", "# last_alpha[b, i, j] is how much position i attended to position j, for example b.\n", "# We print row -1, since that is the row actually used to make the prediction.\n", "ctx, target = next(iter(train_loader))\n", "ctx = ctx.to(attention_model.device)\n", "with torch.no_grad():\n", " attention_model(ctx)\n", "\n", "for b in range(3):\n", " tokens = [tok.decode([i]) for i in ctx[b].tolist()]\n", " weights = attention_model.last_alpha[b, -1].tolist()\n", " print('context: ' + repr(''.join(tokens)) + ' -> next: ' + repr(tok.decode([target[b].item()])))\n", " for t, w in zip(tokens, weights):\n", " print(f' {w:6.3f} {t!r}')\n", " print()" ], "metadata": { "id": "VL_Se52malsS" }, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "source": [ "### From Weighted Averages to a Context-Level MLP\n", "\n", "Attention got the tokens mixing, but notice what it *didn't* do: the weights $\\alpha$ are\n", "just dot products of the token representations with each other, so there are no\n", "parameters in that mixing at all.\n", "And however the weights come out, the result is still a weighted **sum** of per-token\n", "features, which the final predictor then has to work with.\n", "\n", "Let's try removing that pooling bottleneck a different way. \n", "What if we embed each token, concatenate the entire fixed-length context\n", "(`L x d` → `L*d`), and feed this single vector through a multi-layer\n", "perceptron before scoring. \n", "Because the MLP operates on the *whole* context vector, it can freely mix\n", "dimensions from different positions in every layer.\n", "\n", "#### The motivation for this change\n", "* **Broader function class.** \n", " For a fixed context length, a sufficiently wide MLP can approximate any\n", "mapping from the concatenated embeddings to the output, while the positional model was limited to additive “bag-of-tokens” functions.\n", "\n", "* **Position awareness without extra machinery.** \n", " The concatenation preserves a consistent slot order\n", "(position 1 occupies the first block of `d` features, position 2 the next, etc.).\n", "The weight matrices naturally learn different transformations for each slot,\n", "so no separate positional embeddings are required.\n", "\n", "This architecture is still simpler than a Transformer-there is no self-attention\n", "or dynamic token-token interaction but it is a meaningful next step in\n", "expressivity beyond token-wise pooling.\n", "\n", "Task: Complete the TODO: sections in the FeedForwardContextModel class\n" ], "metadata": { "id": "rU-G0yF4E3Cf" } }, { "cell_type": "code", "source": [ "import torch\n", "import torch.nn as nn\n", "\n", "class FeedForwardContextModel(SingleVectorWordModel):\n", " \"\"\"\n", " Context-level MLP (no positional embeddings):\n", " tokens -> embed -> flatten (B, L*d) -> MLP -> score\n", " This model treats the concatenated embeddings of all L tokens\n", " as a single input vector, allowing cross-token interactions.\n", " \"\"\"\n", " def __init__(self,\n", " vocab_size: int,\n", " embedding_dim: int,\n", " ff_hidden_dim: int,\n", " max_context: int):\n", " super().__init__(vocab_size, embedding_dim)\n", " self.max_context = int(max_context)\n", "\n", " # TODO: define MLP: (L*d) -> hidden -> d\n", " self.ff = None\n", "\n", " # TODO: initialize MLP weights (e.g., xavier) and biases to zero\n", " # for m in self.ff:\n", " # if isinstance(m, nn.Linear):\n", " # ...\n", "\n", " def forward(self, context_indices: torch.LongTensor) -> torch.Tensor:\n", " \"\"\"\n", " context_indices: (B, L) with L == max_context\n", " \"\"\"\n", " B, L = context_indices.shape\n", " assert L == self.max_context, f\"Expected context length {self.max_context}, got {L}\"\n", "\n", " # TODO: 1) Embed tokens -> (B, L, d)\n", " x = None\n", "\n", " # TODO: 2) Flatten entire context -> (B, L*d)\n", " x_flat = None\n", "\n", " # TODO: 3) Context-level MLP -> (B, d)\n", " h = None\n", "\n", " # TODO: 4) Score with tied vocabulary embeddings -> (B, V)\n", " logits = None\n", "\n", " return logits\n" ], "metadata": { "id": "SFZbuLOcE2ay" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "HIDDEN_DIM = 1028\n", "mlp = FeedForwardContextModel(VOCAB_SIZE, EMBEDDING_DIM, HIDDEN_DIM, 10)\n", "mlp.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "mlp.to(mlp.device)\n", "\n", "print('Model has {} parameters'.format(sum([np.prod(p.size()) for p in mlp.parameters()])))\n", "\n", "history = train_model(mlp, train_loader, learning_rate, tokenizer=tok, max_steps=10000, experiment_name = \"FeedForwardContextModel1\", val_loader=val_loader)\n" ], "metadata": { "id": "68wPZUBZsRI9" }, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "source": [ "### Going Deeper\n", "\n", "The current context level mlp applies a **single feed-forward block**\n", "to the flattened `L*d` context vector before prediction. \n", "A natural next question is: *how does the model’s **expressive power** change if\n", "we stack multiple such blocks?*\n", "\n", "In function-approximation terms, a one-layer MLP represents functions of the form\n", "\n", "$$\n", "f(x) = \\sigma(Wx + b)\n", "$$\n", "\n", "where \\(\\sigma\\) is a fixed non-linearity. \n", "Adding more layers composes these transformations:\n", "\n", "$$\n", "f(x) = W_k \\,\\sigma\\big(W_{k-1}\\,\\sigma(\\dots \\sigma(W_1 x)\\dots )\\big)\n", "$$\n", "\n", "Classical results in approximation theory show that **compositions of simple\n", "nonlinearities can approximate a far richer class of functions** than a single\n", "layer of the same width. \n", "Unlike the positional model, the input here is already the full context vector,\n", "so deeper stacks directly expand the space of context-level mappings the model\n", "can learn.\n", "\n", "\n", "#### Task\n", "1. Add a `num_layers` parameter so the feed-forward block can be repeated\n", " multiple times. \n", " Begin with `num_layers = 1` (the current model) and experiment with `2`, `4`,\n", " and `8`.\n", "2. Compute how the parameter count scales with depth.\n", "3. Train each version under identical hyper-parameters and compare:\n", " * training loss and validation perplexity\n", " * gradient norms in each layer (use `print_grad_stats`) to monitor stability.\n", "\n", "\n", "\n", "#### Things to Observe\n", "* **Expressivity vs. optimization** – deeper networks can capture more complex\n", "context functions but may be harder to train.\n", "* **Gradient flow** – check for vanishing or exploding gradients as depth grows.\n", "* **Parameter efficiency** – assess whether additional layers provide measurable\n", "gains relative to their extra parameters and compute.\n" ], "metadata": { "id": "-LRW7GNHJAYm" } }, { "cell_type": "code", "source": [ "def print_grad_stats(model):\n", " \"\"\"Print L2 norms of gradients for each layer.\"\"\"\n", " for name, p in model.named_parameters():\n", " if p.grad is not None:\n", " print(f\"{name:<30s} grad_norm={p.grad.norm().item():.6f}\")" ], "metadata": { "id": "QvhmaCFu67fm" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "class _ContextMLPBlock(nn.Module):\n", " \"\"\"\n", " Minimal width-preserving block operating on a d-dim vector:\n", " d -> hidden -> d (with an activation in between).\n", " No residuals, no LayerNorm—keeps the step small.\n", " \"\"\"\n", " def __init__(self, d: int, hidden: int, activation=nn.ReLU):\n", " super().__init__()\n", " # TODO: define first linear layer (d -> hidden)\n", " self.fc1 = None\n", " # TODO: activation function\n", " self.act = None\n", " # TODO: define second linear layer (hidden -> d)\n", " self.fc2 = None\n", " # TODO: initialize weights (e.g., xavier) and biases to zero\n", "\n", " def forward(self, x: torch.Tensor) -> torch.Tensor:\n", " # TODO: forward pass through fc1 -> activation -> fc2\n", " return None\n", "\n", "\n", "class FeedForwardContextModelDeep(SingleVectorWordModel):\n", " \"\"\"\n", " Context-level MLP (no positional embeddings):\n", " tokens -> embed -> flatten (B, L*d) -> stem -> [block] * num_layers -> score\n", "\n", " - Unlike the original FeedForwardContextModel, the stem first maps (L*d) -> d,\n", " and num_layers blocks of d -> hidden -> d are applied on top.\n", " - Increasing num_layers deepens only the d→hidden→d mapping.\n", " \"\"\"\n", " def __init__(self,\n", " vocab_size: int,\n", " embedding_dim: int, # d\n", " ff_hidden_dim: int, # hidden width inside each block\n", " max_context: int, # L\n", " num_layers: int = 1,\n", " activation=nn.ReLU):\n", " super().__init__(vocab_size, embedding_dim)\n", " self.max_context = int(max_context)\n", " self.num_layers = int(num_layers)\n", " d = embedding_dim\n", " L = self.max_context\n", "\n", " # TODO: Stem layer: (L*d) -> d\n", " self.stem_fc = None\n", " self.stem_act = None\n", " # TODO: initialize stem weights and biases\n", "\n", " # TODO: create ModuleList of num_layers ContextMLPBlocks\n", " # (pass in all params, e.g. _ContextMLPBlock(d, ff_hidden_dim, activation))\n", " self.blocks = None\n", "\n", " def forward(self, context_indices: torch.LongTensor) -> torch.Tensor:\n", " \"\"\"\n", " context_indices: (B, L) with L == max_context\n", " \"\"\"\n", " B, L = context_indices.shape\n", " assert L == self.max_context, f\"Expected context length {self.max_context}, got {L}\"\n", "\n", " # TODO: 1) Embed tokens -> (B, L, d)\n", " x = None\n", "\n", " # TODO: 2) Flatten entire context -> (B, L*d)\n", " x_flat = None\n", "\n", " # TODO: 3) Stem to d -> (B, d). Remember to apply the stem activation.\n", " h = None\n", "\n", " # TODO: 4) Depth: apply each block (d -> hidden -> d)\n", " # for blk in self.blocks:\n", " # h = blk(h)\n", "\n", " # TODO: 5) Tied scoring -> (B, V)\n", " logits = None\n", "\n", " return logits\n" ], "metadata": { "id": "lI_I8hHC1wqZ" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "#approximately 10 mins to train for 4 layers of hidden_dim = 1028, 15k steps.\n", "HIDDEN_DIM = 1028\n", "mlpdeep = FeedForwardContextModelDeep(VOCAB_SIZE, EMBEDDING_DIM, HIDDEN_DIM, max_context=10, num_layers = 4)\n", "mlpdeep.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "mlpdeep.to(mlpdeep.device)\n", "\n", "print('Model has {} parameters'.format(sum([np.prod(p.size()) for p in mlpdeep.parameters()])))\n", "\n", "history = train_model(mlpdeep, train_loader, learning_rate, tokenizer=tok, max_steps=15000, experiment_name = \"FeedForwardContextModelDeep1\", val_loader=val_loader)\n" ], "metadata": { "id": "GMzSnsO13SBu" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "print_grad_stats(mlpdeep)" ], "metadata": { "id": "6TrZXMrVe-7M" }, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "source": [ "### Residual Connections\n", "\n", "When we stack many fully-connected layers, training can slow down or become\n", "unstable because gradients shrink or grow as they pass through each layer.\n", "A **residual connection** helps by adding the input of a block directly to its\n", "output:\n", "\n", "$$\n", "\\text{Block}(x) = x + f(x)\n", "$$\n", "\n", "where \\(f(x)\\) is the small network inside the block.\n", "\n", "**Why this helps**\n", "* The gradient can flow along the simple identity path \\(x \\rightarrow x\\),\n", " so learning is less sensitive to depth or initialization.\n", "* Each block only needs to learn a *correction* \\(f(x)\\) on top of the input,\n", " instead of the entire mapping from scratch.\n", "\n", "#### Tasks\n", "\n", "1. **Implement a `ResidualBlock`**\n", " * Input and output both have size `d`.\n", " * Inside the block: `Linear(d→h)` → activation (e.g., ReLU) → `Linear(h→d)`\n", " to compute \\(f(x)\\).\n", " * Return `x + f(x)`.\n", "\n", "2. **Build a `ResidualFeedForwardModel`**\n", " * Start from your context-level MLP pipeline\n", " (`embed → flatten → stem` to dimension `d`).\n", " * Stack several `ResidualBlock`s in a `ModuleList`.\n", " * After the stack, use the same `score()` method to produce logits.\n", "\n", "3. **Experiment**\n", " * Train deep models with and without residuals using the same settings.\n", " * Plot training and validation loss to compare convergence speed.\n", " * Use `print_grad_stats` to inspect gradient norms across layers.\n" ], "metadata": { "id": "n1Z7YAOHffhj" } }, { "cell_type": "code", "source": [ "class ResidualBlock(nn.Module):\n", " \"\"\"\n", " Width-preserving residual block on a d-dim vector:\n", " x -> Linear(d->hidden) -> Act -> Linear(hidden->d) -> + x\n", " \"\"\"\n", " def __init__(self, d: int, hidden: int, activation=nn.ReLU, use_layernorm: bool = False):\n", " super().__init__()\n", " self.use_layernorm = use_layernorm\n", " # TODO: optional LayerNorm or Identity\n", " self.ln = None\n", "\n", " # TODO: first linear layer (d -> hidden)\n", " self.fc1 = None\n", " # TODO: activation function\n", " self.act = None\n", " # TODO: second linear layer (hidden -> d)\n", " self.fc2 = None\n", " # TODO: initialize weights (e.g., xavier) and zero biases\n", "\n", " def forward(self, x: torch.Tensor) -> torch.Tensor:\n", " # TODO: optional pre-norm, forward through fc1 -> activation -> fc2, add residual\n", " return None\n", "\n", "\n", "class ResidualFeedForwardContextModel(SingleVectorWordModel):\n", " \"\"\"\n", " Context-level residual MLP:\n", " tokens -> embed -> flatten (B, L*d) -> stem (L*d -> d)\n", " -> [ResidualBlock(d, hidden)] * num_layers\n", " -> score (tied embeddings)\n", " \"\"\"\n", " def __init__(self,\n", " vocab_size: int,\n", " embedding_dim: int, # d\n", " max_context: int, # L\n", " hidden_dim: int, # hidden width inside each residual block\n", " num_layers: int = 2, # number of residual blocks\n", " activation=nn.ReLU,\n", " use_layernorm: bool = False):\n", " super().__init__(vocab_size, embedding_dim)\n", " self.max_context = int(max_context)\n", " d = embedding_dim\n", " L = self.max_context\n", "\n", " # TODO: Stem (L*d -> d) with activation\n", " self.stem = None\n", " # TODO: initialize stem weights/biases\n", "\n", " # TODO: stack of residual blocks\n", " # (pass in all params, e.g. ResidualBlock(d, hidden_dim, activation, use_layernorm))\n", " self.blocks = None\n", "\n", " def forward(self, context_indices: torch.LongTensor) -> torch.Tensor:\n", " \"\"\"\n", " context_indices: (B, L) with L == max_context\n", " \"\"\"\n", " B, L = context_indices.shape\n", " assert L == self.max_context, f\"Expected context length {self.max_context}, got {L}\"\n", "\n", " # TODO: 1) Embed tokens -> (B, L, d)\n", " x = None\n", "\n", " # TODO: 2) Flatten entire context -> (B, L*d)\n", " x_flat = None\n", "\n", " # TODO: 3) Stem to d -> (B, d). Remember to apply the stem activation.\n", " h = None\n", "\n", " # TODO: 4) Apply residual blocks\n", " # for blk in self.blocks:\n", " # h = blk(h)\n", "\n", " # TODO: 5) Tied scoring -> (B, V)\n", " logits = None\n", "\n", " return logits\n" ], "metadata": { "id": "eysVp08SNmtI" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "residualmlp = ResidualFeedForwardContextModel(VOCAB_SIZE, EMBEDDING_DIM, hidden_dim = HIDDEN_DIM, max_context=10, num_layers=4)\n", "residualmlp.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "residualmlp.to(residualmlp.device)\n", "\n", "print('Model has {} parameters'.format(sum([np.prod(p.size()) for p in residualmlp.parameters()])))\n", "\n", "history = train_model(residualmlp, train_loader, learning_rate, tokenizer=tok, max_steps=15000, experiment_name = \"ResidualFeedForwardContextModel1\", val_loader=val_loader)\n" ], "metadata": { "id": "Ykh5g9uZxJfr" }, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "source": [ "### Gating\n", "\n", "Gating is a learned, elementwise mask that modulates a signal. \n", "Given a representation \\(x\\), a gate \\(g(x)\\) produces per-dimension weights in \\([0,1]\\), and the modulated output is\n", "$$\n", "y \\;=\\; \\sigma(g(x)) \\odot h(x)\n", "$$\n", "where \\(h(x)\\) is any transformation of \\(x\\), \\(\\sigma\\) is a sigmoid, and \\(\\odot\\) is elementwise multiply.\n", "\n", "**Why use it**\n", "- **Selective flow:** the model can amplify useful features and suppress noisy ones.\n", "- **Dynamic sparsity:** encourages the network to use only the channels it needs for a given input." ], "metadata": { "id": "-7bo3eCZSRH_" } }, { "cell_type": "code", "source": [ "class GatedResidualBlock(nn.Module):\n", " \"\"\"\n", " Width-preserving gated residual block on a d-dim vector:\n", " y = x + sigmoid(g(x)) * f(x)\n", " where f(x): d -> h -> d, g(x): d -> d\n", " \"\"\"\n", " def __init__(self, d: int, h: int, activation=nn.ReLU, use_layernorm: bool=False):\n", " super().__init__()\n", " # TODO: optional LayerNorm or Identity\n", " self.ln = None\n", "\n", " # TODO: define f1 (d -> h), activation, f2 (h -> d)\n", " self.f1 = None\n", " self.act = None\n", " self.f2 = None\n", "\n", " # TODO: define gating layer g (d -> d)\n", " self.g = None\n", "\n", " # TODO: initialize all weights/biases (e.g., xavier) and set gate bias to start closed\n", "\n", " def forward(self, x: torch.Tensor) -> torch.Tensor:\n", " # TODO: apply layernorm (if used), compute f(x), compute gate(x),\n", " # and return x + sigmoid(gate) * f(x)\n", " return None\n", "\n", "\n", "class GatedResidualFeedForwardContextModel(SingleVectorWordModel):\n", " \"\"\"\n", " Context-level gated residual MLP (no positional embeddings):\n", " tokens -> embed -> flatten (B, L*d) -> stem (L*d -> d)\n", " -> [GatedResidualBlock(d,h)] * num_layers\n", " -> score() (tied embeddings from SingleVectorWordModel)\n", " \"\"\"\n", " def __init__(self,\n", " vocab_size: int,\n", " embedding_dim: int, # d\n", " max_context: int, # L\n", " hidden_dim: int, # h inside each block\n", " num_layers: int = 2,\n", " activation=nn.ReLU,\n", " use_layernorm: bool=False):\n", " super().__init__(vocab_size, embedding_dim)\n", " self.max_context = int(max_context)\n", " d, L = embedding_dim, self.max_context\n", "\n", " # TODO: Stem layer (L*d -> d) with activation\n", " self.stem = None\n", " # TODO: initialize stem weights/biases\n", "\n", " # TODO: stack of gated residual blocks\n", " # (pass in all params, e.g. GatedResidualBlock(d, hidden_dim, activation, use_layernorm))\n", " self.blocks = None\n", "\n", " def forward(self, context_indices: torch.LongTensor) -> torch.Tensor:\n", " \"\"\"\n", " context_indices: (B, L) with L == max_context\n", " \"\"\"\n", " B, L = context_indices.shape\n", " assert L == self.max_context, f\"Expected context length {self.max_context}, got {L}\"\n", "\n", " # TODO: 1) Embed tokens -> (B, L, d)\n", " x = None\n", "\n", " # TODO: 2) Flatten -> (B, L*d)\n", " x_flat = None\n", "\n", " # TODO: 3) Stem -> (B, d). Remember to apply the stem activation.\n", " h = None\n", "\n", " # TODO: 4) Apply each gated residual block\n", " # for blk in self.blocks:\n", " # h = blk(h)\n", "\n", " # TODO: 5) Tied scoring -> (B, V)\n", " logits = None\n", "\n", " return logits\n" ], "metadata": { "id": "9usvzp_rRNwz" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "gatedresidualmlp = GatedResidualFeedForwardContextModel(VOCAB_SIZE, EMBEDDING_DIM, hidden_dim = HIDDEN_DIM, num_layers=4, max_context=10)\n", "gatedresidualmlp.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "gatedresidualmlp.to(gatedresidualmlp.device)\n", "\n", "print('Model has {} parameters'.format(sum([np.prod(p.size()) for p in gatedresidualmlp.parameters()])))\n", "\n", "history = train_model(gatedresidualmlp, train_loader, learning_rate, tokenizer=tok, max_steps=15000, experiment_name = \"GatedResidualFeedForwardContextModel1\", val_loader=val_loader)" ], "metadata": { "id": "0W79DHdeaKws" }, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "source": [ "# **Take a look at how the models compare!**" ], "metadata": { "id": "iv1cV1qDz5h3" } }, { "cell_type": "code", "source": [ "%load_ext tensorboard\n", "%tensorboard --logdir runs" ], "metadata": { "id": "PoccIOW7qcHL" }, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "source": [ "Discuss the results that you gathered throughout this assignment in two short paragraphs. You might consider topics such as:\n", "\n", "1. How the different model architectures performed relative to each other.\n", "\n", "2. Trends you observed in training (loss curves, convergence speed, overfitting, etc.).\n", "\n", "3. The impact of hyper-parameters.\n", "\n", "4. Any unexpected behaviors or challenges you encountered.\n", "\n", "Support your discussion with specific observations or metrics where possible." ], "metadata": { "id": "mPjb57AR0KFy" } }, { "cell_type": "markdown", "source": [ "Response here:\n" ], "metadata": { "id": "ppMpkRf50wqy" } } ], "metadata": { "colab": { "provenance": [], "gpuType": "T4" }, "kernelspec": { "display_name": "Python 3", "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.7" }, "accelerator": "GPU" }, "nbformat": 4, "nbformat_minor": 0 }