{ "cells": [ { "cell_type": "markdown", "id": "b5978d9d", "metadata": {}, "source": [ "# \ud83d\ude80 Essential ML Coding Interview Guide for MLEs\n", "### From-Scratch Implementations in NumPy & PyTorch\n", "\n", "This notebook is an end-to-end, interview-focused masterclass covering the **most commonly asked Machine Learning coding questions** in top-tier Machine Learning Engineer (MLE) and Applied Scientist interviews (Meta, OpenAI, Anthropic, Apple, Amazon, Databricks).\n", "\n", "In these interviews, candidates are evaluated on:\n", "1. **Mathematical rigor**: Deriving gradients from first principles.\n", "2. **Vectorization**: Avoiding nested loops and mastering tensor broadcasting.\n", "3. **Numerical stability**: Preventing underflow/overflow (log-sum-exp, stable sigmoid, epsilon offsets).\n", "4. **Deep learning architectures**: Building core building blocks (`nn.Module`, parameter initialization traps, tensor shape transformations).\n", "\n", "---\n", "\n", "## \ud83d\udcd1 Table of Contents\n", "1. [Linear Regression (Gradient Descent & Normal Equation)](#1-linear-regression)\n", "2. [Logistic Regression (Gradient Descent & Numerical Stability)](#2-logistic-regression)\n", "3. [K-Means Clustering (Random Choice & List Comprehension)](#3-k-means-clustering)\n", "4. [K-Nearest Neighbors (KNN - Unified Distance & Top-K)](#4-k-nearest-neighbors-knn)\n", "5. [2-Layer MLP with Backpropagation from Scratch (NumPy)](#5-2-layer-mlp-with-backpropagation-from-scratch)\n", "6. [Layer Normalization (LayerNorm in NumPy & PyTorch)](#6-layer-normalization-layernorm)\n", "7. [Multi-Head Self-Attention (Transformer Core)](#7-multi-head-self-attention)\n", "8. [Low-Rank Adaptation (LoRA) Layer](#8-low-rank-adaptation-lora)\n", "9. [Byte-Pair Encoding (BPE) Tokenizer](#9-byte-pair-encoding-bpe)\n", "10. [MLE Interview Quick Reference & Cheat Sheet](#10-mle-interview-quick-reference--cheat-sheet)\n" ] }, { "cell_type": "markdown", "id": "953fd98c", "metadata": {}, "source": [ "---\n", "## 1. Linear Regression\n", "\n", "Linear Regression models a linear relationship between input features $X \\in \\mathbb{R}^{N \\times D}$ and continuous targets $y \\in \\mathbb{R}^{N}$.\n", "\n", "### Mathematical Formulation\n", "$$\n", "\\hat{y} = X w + b\n", "$$\n", "\n", "We optimize the Mean Squared Error (MSE) objective:\n", "$$\n", "\\mathcal{L}(w, b) = \\frac{1}{2N} \\sum_{i=1}^N (\\hat{y}_i - y_i)^2 = \\frac{1}{2N} \\|X w + b - y\\|_2^2\n", "$$\n", "\n", "*(The $\\frac{1}{2}$ prefactor simplifies the derivative by canceling with the exponent 2).*\n", "\n", "### Gradients Derivation\n", "Taking partial derivatives with respect to weights $w$ and bias $b$:\n", "$$\n", "\\frac{\\partial \\mathcal{L}}{\\partial w} = \\frac{1}{N} X^T (\\hat{y} - y) \\in \\mathbb{R}^{D}\n", "$$\n", "$$\n", "\\frac{\\partial \\mathcal{L}}{\\partial b} = \\frac{1}{N} \\sum_{i=1}^N (\\hat{y}_i - y_i) = \\frac{1}{N} \\mathbf{1}^T (\\hat{y} - y) \\in \\mathbb{R}\n", "$$\n", "\n", "### Closed-Form Analytical Solution (Normal Equation)\n", "By augmenting $X$ with an intercept column of ones $\\tilde{X} = [X, \\mathbf{1}] \\in \\mathbb{R}^{N \\times (D+1)}$ and setting $\\nabla_{\\tilde{w}} \\mathcal{L} = 0$:\n", "$$\n", "\\tilde{w} = (\\tilde{X}^T \\tilde{X})^{-1} \\tilde{X}^T y\n", "$$\n", "\n", "> **\ud83d\udca1 Interview Gotchas & Pro-Tips:**\n", "> - **Invertibility trap:** If features are collinear (or $D > N$), $\\tilde{X}^T \\tilde{X}$ is singular (non-invertible). Always mention or use the Moore-Penrose pseudo-inverse (`np.linalg.pinv`), which computes the minimum-norm solution via SVD.\n", "> - **Feature Scaling:** Gradient descent requires feature normalization ($z$-score standardizing $X$) to ensure spherical loss contours and fast convergence. The Normal Equation does not require feature scaling.\n", "> - **Complexity:** Normal Equation is $\\mathcal{O}(D^3 + N D^2)$ time and $\\mathcal{O}(D^2)$ memory. Gradient Descent is $\\mathcal{O}(\\text{epochs} \\cdot N D)$. For large $D$ ($D > 10^4$), GD is strongly preferred.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "ea851c83", "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "\n", "class LinearRegressionGD:\n", " \"\"\"\n", " Linear Regression implemented from scratch with NumPy.\n", " Supports both iterative Batch Gradient Descent and the analytical Normal Equation.\n", " \"\"\"\n", " def __init__(self, lr: float = 0.01, n_iters: int = 1000):\n", " self.lr = lr\n", " self.n_iters = n_iters\n", " self.weights = None\n", " self.bias = None\n", " self.loss_history = []\n", "\n", " def fit(self, X: np.ndarray, y: np.ndarray) -> \"LinearRegressionGD\":\n", " \"\"\"\n", " Fit weights using vectorized batch gradient descent.\n", " X: shape (N, D)\n", " y: shape (N,) or (N, 1)\n", " \"\"\"\n", " N, D = X.shape\n", " y = y.reshape(-1, 1)\n", " \n", " # Initialize parameters: zeros\n", " self.weights = np.zeros((D, 1))\n", " self.bias = 0.0\n", " self.loss_history = []\n", "\n", " for _ in range(self.n_iters):\n", " # 1. Forward prediction\n", " y_pred = np.dot(X, self.weights) + self.bias # Shape: (N, 1)\n", " \n", " # 2. Compute error and MSE loss\n", " error = y_pred - y # Shape: (N, 1)\n", " loss = (1.0 / (2.0 * N)) * np.sum(error ** 2)\n", " self.loss_history.append(loss)\n", " \n", " # 3. Vectorized gradients\n", " dw = (1.0 / N) * np.dot(X.T, error) # Shape: (D, 1)\n", " db = (1.0 / N) * np.sum(error) # Scalar\n", " \n", " # 4. Parameter update\n", " self.weights -= self.lr * dw\n", " self.bias -= self.lr * db\n", " \n", " return self\n", "\n", " def predict(self, X: np.ndarray) -> np.ndarray:\n", " \"\"\"Predict continuous values for input X.\"\"\"\n", " return (np.dot(X, self.weights) + self.bias).squeeze()\n", "\n", " def fit_normal_equation(self, X: np.ndarray, y: np.ndarray) -> \"LinearRegressionGD\":\n", " \"\"\"\n", " Analytical solution: w_aug = (X_aug^T X_aug)^(-1) X_aug^T y\n", " Uses pseudo-inverse for numerical stability with singular matrices.\n", " \"\"\"\n", " N, D = X.shape\n", " y = y.reshape(-1, 1)\n", " \n", " # Augment X with a column of 1s for the bias term\n", " X_aug = np.hstack([X, np.ones((N, 1))]) # Shape: (N, D + 1)\n", " \n", " # Compute pseudo-inverse to handle collinearity gracefully\n", " w_aug = np.linalg.pinv(X_aug.T @ X_aug) @ (X_aug.T @ y)\n", " \n", " self.weights = w_aug[:-1]\n", " self.bias = float(w_aug[-1, 0])\n", " return self\n" ] }, { "cell_type": "code", "execution_count": null, "id": "fed473fc", "metadata": {}, "outputs": [], "source": [ "# --- Verification & Unit Test: Linear Regression ---\n", "np.random.seed(42)\n", "N, D = 200, 3\n", "\n", "# True parameters: y = 3.5*x0 - 2.0*x1 + 1.2*x2 + 4.0 + noise\n", "true_w = np.array([[3.5], [-2.0], [1.2]])\n", "true_b = 4.0\n", "\n", "X = np.random.randn(N, D)\n", "noise = np.random.randn(N, 1) * 0.05\n", "y = X @ true_w + true_b + noise\n", "\n", "# 1. Fit with Gradient Descent\n", "model_gd = LinearRegressionGD(lr=0.05, n_iters=1500)\n", "model_gd.fit(X, y)\n", "\n", "# 2. Fit with Normal Equation\n", "model_ne = LinearRegressionGD()\n", "model_ne.fit_normal_equation(X, y)\n", "\n", "print(\"Linear Regression Results:\")\n", "print(f\"True Weights: {true_w.ravel()}, True Bias: {true_b:.4f}\")\n", "print(f\"GD Weights: {model_gd.weights.ravel()}, GD Bias: {model_gd.bias:.4f}\")\n", "print(f\"NE Weights: {model_ne.weights.ravel()}, NE Bias: {model_ne.bias:.4f}\")\n", "\n", "# Check that GD converged close to Normal Equation\n", "assert np.allclose(model_gd.weights, model_ne.weights, atol=1e-2), \"GD weights did not match Normal Equation!\"\n", "assert np.isclose(model_gd.bias, model_ne.bias, atol=1e-2), \"GD bias did not match Normal Equation!\"\n", "print(\"\u2705 Linear Regression passed all verification checks!\")\n" ] }, { "cell_type": "markdown", "id": "ef61373a", "metadata": {}, "source": [ "---\n", "## 2. Logistic Regression\n", "\n", "Logistic Regression is the foundational classification model that maps linear combinations of features to class probabilities $\\hat{y} \\in (0, 1)$ via the Sigmoid activation function.\n", "\n", "### Mathematical Formulation\n", "$$\n", "\\hat{y} = \\sigma(z) = \\frac{1}{1 + e^{-z}}, \\quad \\text{where } z = X w + b\n", "$$\n", "\n", "The objective is to minimize the **Binary Cross-Entropy (BCE) Loss**:\n", "$$\n", "\\mathcal{L}(w, b) = -\\frac{1}{N} \\sum_{i=1}^N \\left[ y_i \\log(\\hat{y}_i) + (1 - y_i) \\log(1 - \\hat{y}_i) \\right]\n", "$$\n", "\n", "### Gradients Derivation\n", "Using the chain rule and the identity $\\frac{d\\sigma(z)}{dz} = \\sigma(z)(1 - \\sigma(z))$:\n", "$$\n", "\\frac{\\partial \\mathcal{L}}{\\partial \\hat{y}_i} = -\\frac{y_i}{\\hat{y}_i} + \\frac{1 - y_i}{1 - \\hat{y}_i} = \\frac{\\hat{y}_i - y_i}{\\hat{y}_i(1 - \\hat{y}_i)}\n", "$$\n", "$$\n", "\\frac{\\partial \\hat{y}_i}{\\partial z_i} = \\hat{y}_i(1 - \\hat{y}_i)\n", "$$\n", "Multiplying gives:\n", "$$\n", "\\frac{\\partial \\mathcal{L}}{\\partial z_i} = \\hat{y}_i - y_i\n", "$$\n", "Vectorized across all samples:\n", "$$\n", "\\frac{\\partial \\mathcal{L}}{\\partial w} = \\frac{1}{N} X^T (\\hat{y} - y) \\in \\mathbb{R}^D\n", "$$\n", "$$\n", "\\frac{\\partial \\mathcal{L}}{\\partial b} = \\frac{1}{N} \\sum_{i=1}^N (\\hat{y}_i - y_i) \\in \\mathbb{R}\n", "$$\n", "\n", "> **\ud83d\udca1 Interview Gotchas & Pro-Tips:**\n", "> - **Sigmoid Overflow / Underflow:** When $z \\ll 0$, $e^{-z} \\to \\infty$ which raises `OverflowError`. A numerically stable sigmoid branches based on sign:\n", "> $$\n", "> \\sigma(z) = \\begin{cases} \\frac{1}{1 + e^{-z}} & \\text{if } z \\ge 0 \\\\ \\frac{e^z}{1 + e^z} & \\text{if } z < 0 \\end{cases}\n", "> $$\n", "> - **Log(0) Explosion:** If $\\hat{y} = 0$, $\\log(\\hat{y}) = -\\infty$. Always clip probabilities: `np.clip(y_pred, 1e-15, 1 - 1e-15)`.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "0e884689", "metadata": {}, "outputs": [], "source": [ "class LogisticRegressionGD:\n", " \"\"\"\n", " Logistic Regression with numerical stability tricks implemented from scratch in NumPy.\n", " \"\"\"\n", " def __init__(self, lr: float = 0.1, n_iters: int = 1000):\n", " self.lr = lr\n", " self.n_iters = n_iters\n", " self.weights = None\n", " self.bias = None\n", " self.loss_history = []\n", "\n", " @staticmethod\n", " def _stable_sigmoid(z: np.ndarray) -> np.ndarray:\n", " \"\"\"\n", " Numerically stable sigmoid function:\n", " For z >= 0: 1 / (1 + exp(-z))\n", " For z < 0: exp(z) / (1 + exp(z))\n", " \"\"\"\n", " z = np.asarray(z, dtype=np.float64)\n", " positive_mask = (z >= 0)\n", " negative_mask = ~positive_mask\n", " \n", " result = np.empty_like(z)\n", " result[positive_mask] = 1.0 / (1.0 + np.exp(-z[positive_mask]))\n", " exp_z_neg = np.exp(z[negative_mask])\n", " result[negative_mask] = exp_z_neg / (1.0 + exp_z_neg)\n", " return result\n", "\n", " def fit(self, X: np.ndarray, y: np.ndarray) -> \"LogisticRegressionGD\":\n", " N, D = X.shape\n", " y = y.reshape(-1, 1)\n", " \n", " self.weights = np.zeros((D, 1))\n", " self.bias = 0.0\n", " self.loss_history = []\n", "\n", " eps = 1e-15 # For stable log computation\n", "\n", " for _ in range(self.n_iters):\n", " # 1. Forward linear combination + sigmoid\n", " linear_output = np.dot(X, self.weights) + self.bias\n", " y_pred = self._stable_sigmoid(linear_output)\n", " \n", " # 2. Binary Cross-Entropy loss with numerical clipping\n", " y_pred_clipped = np.clip(y_pred, eps, 1.0 - eps)\n", " loss = - (1.0 / N) * np.sum(\n", " y * np.log(y_pred_clipped) + (1.0 - y) * np.log(1.0 - y_pred_clipped)\n", " )\n", " self.loss_history.append(loss)\n", " \n", " # 3. Vectorized gradients\n", " error = y_pred - y\n", " dw = (1.0 / N) * np.dot(X.T, error)\n", " db = (1.0 / N) * np.sum(error)\n", " \n", " # 4. Gradient descent update\n", " self.weights -= self.lr * dw\n", " self.bias -= self.lr * db\n", " \n", " return self\n", "\n", " def predict_proba(self, X: np.ndarray) -> np.ndarray:\n", " \"\"\"Return predicted probability for class 1.\"\"\"\n", " linear_output = np.dot(X, self.weights) + self.bias\n", " return self._stable_sigmoid(linear_output).squeeze()\n", "\n", " def predict(self, X: np.ndarray, threshold: float = 0.5) -> np.ndarray:\n", " \"\"\"Return binary class label (0 or 1).\"\"\"\n", " return (self.predict_proba(X) >= threshold).astype(int)\n" ] }, { "cell_type": "code", "execution_count": null, "id": "ef62c8fe", "metadata": {}, "outputs": [], "source": [ "# --- Verification & Unit Test: Logistic Regression ---\n", "np.random.seed(42)\n", "N = 300\n", "\n", "# Create two clearly separable 2D Gaussian clusters\n", "X0 = np.random.randn(N // 2, 2) + np.array([-2.0, -2.0])\n", "X1 = np.random.randn(N // 2, 2) + np.array([2.0, 2.0])\n", "X = np.vstack([X0, X1])\n", "y = np.array([0] * (N // 2) + [1] * (N // 2))\n", "\n", "# Train model\n", "clf = LogisticRegressionGD(lr=0.1, n_iters=1000)\n", "clf.fit(X, y)\n", "\n", "preds = clf.predict(X)\n", "accuracy = np.mean(preds == y)\n", "\n", "print(f\"Trained Logistic Regression Accuracy: {accuracy * 100:.2f}%\")\n", "print(f\"Initial Loss: {clf.loss_history[0]:.4f} -> Final Loss: {clf.loss_history[-1]:.4f}\")\n", "\n", "assert accuracy > 0.95, f\"Expected >95% accuracy on separable data, got {accuracy * 100}%\"\n", "assert clf.loss_history[-1] < clf.loss_history[0], \"Loss failed to decrease during training!\"\n", "print(\"\u2705 Logistic Regression passed all verification checks!\")\n" ] }, { "cell_type": "markdown", "id": "a20549c0", "metadata": {}, "source": [ "---\n", "## 3. K-Means Clustering\n", "\n", "K-Means partitions $N$ observations into $K$ clusters such that each observation belongs to the cluster with the nearest mean (centroid).\n", "\n", "### Mathematical Objective\n", "Minimize within-cluster sum of squared distances:\n", "$$\n", "\\mathcal{J} = \\sum_{k=1}^K \\sum_{x \\in \\mathcal{C}_k} \\|x - \\mu_k\\|_2^2\n", "$$\n", "\n", "### Algorithm Steps in a Single Unified Function\n", "1. **Centroid Initialization:** Randomly choose $K$ distinct data points from $X$ without replacement:\n", " $$\n", " \\text{indices} = \\text{random\\_choice}(N, K, \\text{replace}=\\text{False})\n", " $$\n", "2. **Assignment Step (Broadcasting):** Assign each sample $x_i$ to its nearest centroid:\n", " $$\n", " c_i = \\arg\\min_{k \\in \\{1, \\dots, K\\}} \\|x_i - \\mu_k\\|_2^2\n", " $$\n", "3. **Update Step (List Comprehension):** Recalculate centroids as the mean of assigned points using a Python list comprehension:\n", " $$\n", " \\mu_k = \\frac{1}{|\\mathcal{C}_k|} \\sum_{x \\in \\mathcal{C}_k} x\n", " $$\n", "4. **Convergence:** Terminate when the maximum centroid shift across iterations is less than tolerance $\\text{tol}$ or when $\\text{max\\_iters}$ is reached.\n", "\n", "> **\ud83d\udca1 Interview Gotchas & Pro-Tips:**\n", "> - **Unified Function:** Keep initialization and the iteration loop concise within a single `fit()` function.\n", "> - **Pythonic List Comprehension:** Interviewers appreciate concise vectorization: updating centroids with `np.array([np.mean(X[labels == k], axis=0) for k in range(self.k)])`.\n", "> - **Vectorized Broadcasting:** Pairwise distances are computed in one line `np.linalg.norm(X[:, None, :] - self.centroids[None, :, :], axis=2)` without nested loops.\n", "" ] }, { "cell_type": "code", "execution_count": null, "id": "4617ddd1", "metadata": {}, "outputs": [], "source": [ "class KMeans:\n", " \"\"\"\n", " K-Means clustering implemented with a single unified fit() function.\n", " Uses random choice initialization, vectorized distance broadcasting,\n", " and list comprehension for centroid updates.\n", " \"\"\"\n", " def __init__(self, k: int = 3, max_iters: int = 300, tol: float = 1e-4):\n", " self.k = k\n", " self.max_iters = max_iters\n", " self.tol = tol\n", " self.centroids = None\n", "\n", " def fit(self, X: np.ndarray) -> \"KMeans\":\n", " N, D = X.shape\n", "\n", " # 1. Random centroid initialization without replacement\n", " random_indices = np.random.choice(N, size=self.k, replace=False)\n", " self.centroids = X[random_indices].astype(np.float64)\n", "\n", " # 2. Iteration Loop\n", " for _ in range(self.max_iters):\n", " # Vectorized distance computation (broadcasting): shape (N, K)\n", " distances = np.linalg.norm(X[:, None, :] - self.centroids[None, :, :], axis=2)\n", " labels = np.argmin(distances, axis=1)\n", "\n", " # Centroid update using list comprehension\n", " new_centroids = np.array([\n", " np.mean(X[labels == k_idx], axis=0) if np.any(labels == k_idx) else self.centroids[k_idx]\n", " for k_idx in range(self.k)\n", " ])\n", "\n", " # Convergence check: maximum centroid shift < tol\n", " shift = np.max(np.linalg.norm(new_centroids - self.centroids, axis=1))\n", " self.centroids = new_centroids\n", " if shift < self.tol:\n", " break\n", "\n", " return self\n", "\n", " def predict(self, X: np.ndarray) -> np.ndarray:\n", " \"\"\"Assign samples to closest cluster centroid.\"\"\"\n", " distances = np.linalg.norm(X[:, None, :] - self.centroids[None, :, :], axis=2)\n", " return np.argmin(distances, axis=1)\n", "" ] }, { "cell_type": "code", "execution_count": null, "id": "1281a2a2", "metadata": {}, "outputs": [], "source": [ "# --- Verification & Unit Test: K-Means Clustering ---\n", "np.random.seed(42)\n", "\n", "# Generate 3 distinct 2D clusters with 100 points each\n", "c1 = np.random.randn(100, 2) * 0.5 + np.array([-5.0, -5.0])\n", "c2 = np.random.randn(100, 2) * 0.5 + np.array([0.0, 5.0])\n", "c3 = np.random.randn(100, 2) * 0.5 + np.array([5.0, -5.0])\n", "X_clusters = np.vstack([c1, c2, c3])\n", "\n", "kmeans = KMeans(k=3)\n", "kmeans.fit(X_clusters)\n", "\n", "labels = kmeans.predict(X_clusters)\n", "\n", "print(f\"K-Means fitted centroids:\\n{kmeans.centroids}\")\n", "\n", "unique, counts = np.unique(labels, return_counts=True)\n", "print(f\"Cluster assignments count: {dict(zip(unique, counts))}\")\n", "\n", "assert len(unique) == 3, f\"Expected 3 clusters, but found {len(unique)}\"\n", "assert all(85 <= c <= 115 for c in counts), \"Cluster counts indicate poor convergence!\"\n", "print(\"\u2705 K-Means passed all verification checks!\")\n", "" ] }, { "cell_type": "markdown", "id": "543d4c48", "metadata": {}, "source": [ "---\n", "## 4. K-Nearest Neighbors (KNN)\n", "\n", "KNN is a non-parametric, instance-based learning algorithm used for classification and regression.\n", "\n", "### Mathematical Formulation\n", "Given test query $x$ and training points $\\{x_i, y_i\\}_{i=1}^N$:\n", "1. Compute Euclidean distance:\n", " $$\n", " d(x, x_i) = \\|x - x_i\\|_2 = \\sqrt{\\sum_{j=1}^D (x_j - x_{i,j})^2}\n", " $$\n", "2. Identify indices of $K$ smallest distances: $\\mathcal{N}_K(x)$.\n", "3. **Classification:** Output majority vote $\\hat{y} = \\text{mode}(\\{y_i : i \\in \\mathcal{N}_K(x)\\})$.\n", "4. **Regression:** Output mean $\\hat{y} = \\frac{1}{K} \\sum_{i \\in \\mathcal{N}_K(x)} y_i$.\n", "\n", "> **\ud83d\udca1 Interview Gotchas & Pro-Tips:**\n", "> - **Unified Predict Function:** Distance computation, $O(N)$ top-k partition, and voting are executed inside a single `predict()` function without separate helper methods.\n", "> - **Top-K Complexity:**\n", "> - `np.argsort(distances)` takes $\\mathcal{O}(N \\log N)$.\n", "> - `np.argpartition(distances, kth=k)` takes $\\mathcal{O}(N)$ average time. Always use `argpartition` for optimal performance.\n", "> - **Vectorized Broadcasting:** Pairwise distances between $M$ test points and $N$ training points are computed in one line:\n", "> `distances = np.linalg.norm(X_test[:, None, :] - self.X_train[None, :, :], axis=2)`\n" ] }, { "cell_type": "code", "execution_count": null, "id": "0b3275ed", "metadata": {}, "outputs": [], "source": [ "class KNearestNeighbors:\n", " \"\"\"\n", " Vectorized K-Nearest Neighbors in a clean, self-contained implementation.\n", " Computes pairwise distances, O(N) top-K partition, and voting inside predict().\n", " \"\"\"\n", " def __init__(self, k: int = 5, task: str = \"classification\"):\n", " assert task in (\"classification\", \"regression\"), \"task must be 'classification' or 'regression'\"\n", " self.k = k\n", " self.task = task\n", " self.X_train = None\n", " self.y_train = None\n", "\n", " def fit(self, X: np.ndarray, y: np.ndarray) -> \"KNearestNeighbors\":\n", " \"\"\"Lazy learner: store training data.\"\"\"\n", " self.X_train = np.asarray(X, dtype=np.float64)\n", " self.y_train = np.asarray(y)\n", " return self\n", "\n", " def predict(self, X_test: np.ndarray) -> np.ndarray:\n", " \"\"\"\n", " Unified prediction function:\n", " 1. Computes (M, N) pairwise distance matrix via broadcasting.\n", " 2. Retrieves top-k nearest neighbors via O(N) argpartition.\n", " 3. Computes mode (classification) or mean (regression).\n", " \"\"\"\n", " X_test = np.asarray(X_test, dtype=np.float64)\n", " M = X_test.shape[0]\n", "\n", " # Vectorized pairwise Euclidean distances: (M, 1, D) - (1, N, D) -> (M, N)\n", " distances = np.linalg.norm(X_test[:, None, :] - self.X_train[None, :, :], axis=2)\n", "\n", " predictions = []\n", " for i in range(M):\n", " row_dists = distances[i]\n", " # O(N) selection of k nearest neighbor indices\n", " k_indices = np.argpartition(row_dists, kth=self.k - 1)[:self.k]\n", " k_labels = self.y_train[k_indices]\n", "\n", " if self.task == \"classification\":\n", " values, counts = np.unique(k_labels, return_counts=True)\n", " pred = values[np.argmax(counts)]\n", " else:\n", " pred = np.mean(k_labels)\n", "\n", " predictions.append(pred)\n", "\n", " return np.array(predictions)\n" ] }, { "cell_type": "code", "execution_count": null, "id": "984d6d03", "metadata": {}, "outputs": [], "source": [ "# --- Verification & Unit Test: KNN ---\n", "np.random.seed(42)\n", "N_train = 150\n", "\n", "# Generate 3 training classes\n", "X_tr = np.vstack([\n", " np.random.randn(N_train // 3, 2) + np.array([-3.0, 0.0]),\n", " np.random.randn(N_train // 3, 2) + np.array([3.0, 0.0]),\n", " np.random.randn(N_train // 3, 2) + np.array([0.0, 4.0])\n", "])\n", "y_tr = np.array([0] * (N_train // 3) + [1] * (N_train // 3) + [2] * (N_train // 3))\n", "\n", "# Query test points close to each center\n", "X_te = np.array([[-3.0, 0.0], [3.0, 0.0], [0.0, 4.0]])\n", "expected_y = np.array([0, 1, 2])\n", "\n", "knn = KNearestNeighbors(k=5, task=\"classification\")\n", "knn.fit(X_tr, y_tr)\n", "preds = knn.predict(X_te)\n", "\n", "print(f\"KNN Test predictions: {preds}, Expected: {expected_y}\")\n", "assert np.array_equal(preds, expected_y), f\"Expected {expected_y}, got {preds}\"\n", "print(\"\u2705 K-Nearest Neighbors passed all verification checks!\")\n" ] }, { "cell_type": "markdown", "id": "c2189078", "metadata": {}, "source": [ "---\n", "## 5. 2-Layer Multi-Layer Perceptron (Backprop from Scratch in NumPy)\n", "\n", "Implementing an MLP from scratch using pure NumPy proves complete mastery over the **multivariate chain rule**, activation derivatives, and tensor shapes.\n", "\n", "### Architecture & Forward Pass\n", "- **Input:** $X \\in \\mathbb{R}^{N \\times D_{\\text{in}}}$\n", "- **Hidden Layer:**\n", " $$\n", " Z_1 = X W_1 + b_1 \\in \\mathbb{R}^{N \\times H}, \\qquad A_1 = \\text{ReLU}(Z_1) = \\max(0, Z_1)\n", " $$\n", "- **Output Layer:**\n", " $$\n", " Z_2 = A_1 W_2 + b_2 \\in \\mathbb{R}^{N \\times 1}, \\qquad \\hat{Y} = \\sigma(Z_2) \\in \\mathbb{R}^{N \\times 1}\n", " $$\n", "- **Loss:** Binary Cross-Entropy\n", " $$\n", " \\mathcal{L} = -\\frac{1}{N} \\sum_{i=1}^N \\left[ Y_i \\log(\\hat{Y}_i) + (1 - Y_i)\\log(1 - \\hat{Y}_i) \\right]\n", " $$\n", "\n", "### Backward Pass (Analytical Chain Rule Derivation)\n", "1. **Output layer gradient:**\n", " $$\n", " \\delta_2 = \\frac{\\partial \\mathcal{L}}{\\partial Z_2} = \\frac{1}{N} (\\hat{Y} - Y) \\in \\mathbb{R}^{N \\times 1}\n", " $$\n", "2. **Output weights and bias gradients:**\n", " $$\n", " \\frac{\\partial \\mathcal{L}}{\\partial W_2} = A_1^T \\delta_2 \\in \\mathbb{R}^{H \\times 1}, \\qquad \\frac{\\partial \\mathcal{L}}{\\partial b_2} = \\sum_{i=1}^N \\delta_{2, i} \\in \\mathbb{R}^{1 \\times 1}\n", " $$\n", "3. **Hidden layer activation gradient:**\n", " $$\n", " \\frac{\\partial \\mathcal{L}}{\\partial A_1} = \\delta_2 W_2^T \\in \\mathbb{R}^{N \\times H}\n", " $$\n", "4. **Hidden pre-activation gradient:**\n", " $$\n", " \\delta_1 = \\frac{\\partial \\mathcal{L}}{\\partial Z_1} = \\left(\\delta_2 W_2^T\\right) \\odot \\mathbf{1}_{\\{Z_1 > 0\\}} \\in \\mathbb{R}^{N \\times H}\n", " $$\n", "5. **Input weights and bias gradients:**\n", " $$\n", " \\frac{\\partial \\mathcal{L}}{\\partial W_1} = X^T \\delta_1 \\in \\mathbb{R}^{D_{\\text{in}} \\times H}, \\qquad \\frac{\\partial \\mathcal{L}}{\\partial b_1} = \\sum_{i=1}^N \\delta_{1, i} \\in \\mathbb{R}^{1 \\times H}\n", " $$\n", "\n", "> **\ud83d\udca1 Interview Gotchas & Pro-Tips:**\n", "> - **Symmetry Breaking:** Initializing weights to all zeros means every hidden neuron computes identical outputs and receives identical gradients. Always use **He (Kaiming) Normal** initialization: $W \\sim \\mathcal{N}(0, \\sqrt{2 / D_{\\text{in}}})$.\n", "> - **Shape Matching:** Always verify shapes at each stage:\n", "> `X: (N, D) -> Z1: (N, H) -> Z2: (N, 1) -> dW2: (H, 1) -> dW1: (D, H)`.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "235ae36c", "metadata": {}, "outputs": [], "source": [ "class TwoLayerMLP:\n", " \"\"\"\n", " 2-Layer Multi-Layer Perceptron with manual backpropagation implemented in pure NumPy.\n", " Architecture: Input -> Linear -> ReLU -> Linear -> Sigmoid (BCE Loss).\n", " \"\"\"\n", " def __init__(self, in_dim: int, hidden_dim: int, lr: float = 0.1):\n", " self.lr = lr\n", " \n", " # He (Kaiming) Normal Initialization for ReLU\n", " self.W1 = np.random.randn(in_dim, hidden_dim) * np.sqrt(2.0 / in_dim)\n", " self.b1 = np.zeros((1, hidden_dim))\n", " \n", " # Xavier Normal Initialization for Sigmoid output\n", " self.W2 = np.random.randn(hidden_dim, 1) * np.sqrt(1.0 / hidden_dim)\n", " self.b2 = np.zeros((1, 1))\n", " \n", " self.loss_history = []\n", "\n", " @staticmethod\n", " def _sigmoid(z: np.ndarray) -> np.ndarray:\n", " return np.where(z >= 0, 1.0 / (1.0 + np.exp(-z)), np.exp(z) / (1.0 + np.exp(z)))\n", "\n", " def forward(self, X: np.ndarray):\n", " \"\"\"\n", " Forward pass. Returns final predictions and intermediate cache for backprop.\n", " \"\"\"\n", " Z1 = np.dot(X, self.W1) + self.b1 # Shape: (N, H)\n", " A1 = np.maximum(0.0, Z1) # ReLU activation\n", " \n", " Z2 = np.dot(A1, self.W2) + self.b2 # Shape: (N, 1)\n", " A2 = self._sigmoid(Z2) # Sigmoid output\n", " \n", " cache = {\"X\": X, \"Z1\": Z1, \"A1\": A1, \"Z2\": Z2, \"A2\": A2}\n", " return A2, cache\n", "\n", " def backward(self, Y: np.ndarray, cache: dict):\n", " \"\"\"\n", " Backward pass: computes exact analytical gradients via multivariate chain rule.\n", " \"\"\"\n", " N = Y.shape[0]\n", " X = cache[\"X\"]\n", " Z1 = cache[\"Z1\"]\n", " A1 = cache[\"A1\"]\n", " A2 = cache[\"A2\"]\n", " \n", " # 1. Output gradient: dL/dZ2 = (1/N) * (A2 - Y)\n", " dZ2 = (1.0 / N) * (A2 - Y) # Shape: (N, 1)\n", " dW2 = np.dot(A1.T, dZ2) # Shape: (H, 1)\n", " db2 = np.sum(dZ2, axis=0, keepdims=True) # Shape: (1, 1)\n", " \n", " # 2. Backprop to hidden layer: dL/dZ1 = (dZ2 @ W2.T) * ReLU'(Z1)\n", " dA1 = np.dot(dZ2, self.W2.T) # Shape: (N, H)\n", " dZ1 = dA1 * (Z1 > 0).astype(float) # ReLU derivative\n", " \n", " dW1 = np.dot(X.T, dZ1) # Shape: (D, H)\n", " db1 = np.sum(dZ1, axis=0, keepdims=True) # Shape: (1, H)\n", " \n", " # 3. Parameter updates (Gradient Descent)\n", " self.W2 -= self.lr * dW2\n", " self.b2 -= self.lr * db2\n", " self.W1 -= self.lr * dW1\n", " self.b1 -= self.lr * db1\n", "\n", " def fit(self, X: np.ndarray, y: np.ndarray, epochs: int = 2000):\n", " N = X.shape[0]\n", " Y = y.reshape(-1, 1)\n", " eps = 1e-15\n", " \n", " for epoch in range(epochs):\n", " # Forward pass\n", " preds, cache = self.forward(X)\n", " \n", " # Loss computation\n", " preds_clipped = np.clip(preds, eps, 1.0 - eps)\n", " loss = - (1.0 / N) * np.sum(Y * np.log(preds_clipped) + (1.0 - Y) * np.log(1.0 - preds_clipped))\n", " self.loss_history.append(loss)\n", " \n", " # Backward pass\n", " self.backward(Y, cache)\n", " \n", " return self\n", "\n", " def predict(self, X: np.ndarray) -> np.ndarray:\n", " preds, _ = self.forward(X)\n", " return (preds >= 0.5).astype(int).squeeze()\n" ] }, { "cell_type": "code", "execution_count": null, "id": "cdfbc4e9", "metadata": {}, "outputs": [], "source": [ "# --- Verification & Unit Test: 2-Layer MLP on Non-Linear XOR Problem ---\n", "np.random.seed(42)\n", "\n", "# XOR problem: non-linear classification boundary\n", "X_xor = np.array([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=np.float64)\n", "y_xor = np.array([0, 1, 1, 0], dtype=np.float64)\n", "\n", "X_train = np.repeat(X_xor, 50, axis=0) + np.random.randn(200, 2) * 0.05\n", "y_train = np.repeat(y_xor, 50)\n", "\n", "mlp = TwoLayerMLP(in_dim=2, hidden_dim=8, lr=0.5)\n", "mlp.fit(X_train, y_train, epochs=2500)\n", "\n", "xor_preds = mlp.predict(X_xor)\n", "print(\"XOR Truth:\", y_xor)\n", "print(\"XOR Preds:\", xor_preds)\n", "print(f\"MLP Initial Loss: {mlp.loss_history[0]:.4f} -> Final Loss: {mlp.loss_history[-1]:.4f}\")\n", "\n", "assert np.array_equal(xor_preds, y_xor), \"MLP failed to solve non-linear XOR problem!\"\n", "assert mlp.loss_history[-1] < 0.1, \"Loss did not converge!\"\n", "print(\"\u2705 2-Layer MLP passed all verification checks!\")\n" ] }, { "cell_type": "markdown", "id": "b8a187e9", "metadata": {}, "source": [ "---\n", "## 6. Layer Normalization (LayerNorm)\n", "\n", "Introduced by Ba, Kiros, & Hinton (2016), Layer Normalization is an indispensable building block of modern Transformers (GPT, BERT, LLaMA).\n", "\n", "### Mathematical Formulation\n", "Given an input tensor $x \\in \\mathbb{R}^{B \\times T \\times D}$ (or $x \\in \\mathbb{R}^{N \\times D}$):\n", "For each individual sample and token position, we compute mean $\\mu$ and variance $\\sigma^2$ across the **last dimension** (feature dimension $D$):\n", "\n", "$$\n", "\\mu = \\frac{1}{D} \\sum_{j=1}^D x_j, \\qquad \\sigma^2 = \\frac{1}{D} \\sum_{j=1}^D (x_j - \\mu)^2\n", "$$\n", "\n", "Normalize and apply learnable affine parameters $\\gamma$ (scale) and $\\beta$ (shift):\n", "$$\n", "\\hat{x} = \\frac{x - \\mu}{\\sqrt{\\sigma^2 + \\epsilon}}\n", "$$\n", "$$\n", "y = \\gamma \\odot \\hat{x} + \\beta\n", "$$\n", "\n", "Where:\n", "- $\\gamma \\in \\mathbb{R}^D$ is initialized to $\\mathbf{1}$.\n", "- $\\beta \\in \\mathbb{R}^D$ is initialized to $\\mathbf{0}$.\n", "- $\\epsilon > 0$ is a small scalar (e.g. $10^{-5}$) to prevent division by zero.\n", "\n", "---\n", "\n", "### \ud83c\udd9a Critical Interview Comparison: LayerNorm vs. BatchNorm\n", "\n", "| Feature | **BatchNorm** | **LayerNorm** |\n", "|---|---|---|\n", "| **Dimension normalized across** | Across the **batch dimension** $B$ for each feature channel | Across the **feature dimension** $D$ for each sample independently |\n", "| **Dependency between batch samples** | **Yes** (a sample's output depends on other batch members) | **No** (each sample is normalized completely independently) |\n", "| **Behavior during Inference** | Requires tracking running stats (`running_mean`, `running_var`) | Exact same computation during train and inference! |\n", "| **Suitability for Transformers & NLP**| Poor (fails with variable sequence lengths & small batch sizes) | **Ideal** (standard for LLMs and recurrent models) |\n" ] }, { "cell_type": "code", "execution_count": null, "id": "b7e0abf0", "metadata": {}, "outputs": [], "source": [ "# --- Pure NumPy Implementation of LayerNorm ---\n", "def layer_norm_numpy(x: np.ndarray, gamma: np.ndarray, beta: np.ndarray, eps: float = 1e-5) -> np.ndarray:\n", " \"\"\"\n", " x: shape (..., D)\n", " gamma: shape (D,)\n", " beta: shape (D,)\n", " \"\"\"\n", " mean = np.mean(x, axis=-1, keepdims=True)\n", " var = np.var(x, axis=-1, keepdims=True)\n", " x_hat = (x - mean) / np.sqrt(var + eps)\n", " return gamma * x_hat + beta\n", "\n", "# --- PyTorch Module Implementation of LayerNorm ---\n", "try:\n", " import torch\n", " import torch.nn as nn\n", "\n", " class CustomLayerNorm(nn.Module):\n", " def __init__(self, normalized_shape: int, eps: float = 1e-5):\n", " super().__init__()\n", " self.eps = eps\n", " # Learnable parameters: gamma (initialized to 1), beta (initialized to 0)\n", " self.gamma = nn.Parameter(torch.ones(normalized_shape))\n", " self.beta = nn.Parameter(torch.zeros(normalized_shape))\n", "\n", " def forward(self, x: torch.Tensor) -> torch.Tensor:\n", " # x shape: (..., D)\n", " mean = x.mean(dim=-1, keepdim=True)\n", " var = x.var(dim=-1, unbiased=False, keepdim=True) # MLE variance (divide by D, not D-1)\n", " x_norm = (x - mean) / torch.sqrt(var + self.eps)\n", " return self.gamma * x_norm + self.beta\n", "\n", " TORCH_AVAILABLE = True\n", "except ImportError:\n", " TORCH_AVAILABLE = False\n" ] }, { "cell_type": "code", "execution_count": null, "id": "708638fa", "metadata": {}, "outputs": [], "source": [ "# --- Verification: LayerNorm Numerical Accuracy ---\n", "# 1. Test NumPy LayerNorm\n", "x_np = np.array([[1.0, 2.0, 3.0, 4.0], [10.0, 20.0, 30.0, 40.0]])\n", "gamma_np = np.ones(4)\n", "beta_np = np.zeros(4)\n", "\n", "out_np = layer_norm_numpy(x_np, gamma_np, beta_np)\n", "print(\"NumPy LayerNorm output shape:\", out_np.shape)\n", "assert np.allclose(np.mean(out_np, axis=-1), 0.0, atol=1e-6)\n", "assert np.allclose(np.var(out_np, axis=-1), 1.0, atol=1e-3)\n", "print(\"NumPy LayerNorm normalized correctly to mean 0 and var 1.\")\n", "\n", "# 2. Test PyTorch LayerNorm vs official torch.nn.LayerNorm\n", "if TORCH_AVAILABLE:\n", " B, T, D = 2, 4, 8\n", " x_torch = torch.randn(B, T, D)\n", " \n", " custom_ln = CustomLayerNorm(D)\n", " official_ln = nn.LayerNorm(D, eps=1e-5)\n", " \n", " with torch.no_grad():\n", " official_ln.weight.copy_(custom_ln.gamma)\n", " official_ln.bias.copy_(custom_ln.beta)\n", " \n", " custom_out = custom_ln(x_torch)\n", " official_out = official_ln(x_torch)\n", " \n", " assert torch.allclose(custom_out, official_out, atol=1e-5), \"Custom LayerNorm did not match official PyTorch!\"\n", " print(\"\u2705 PyTorch CustomLayerNorm matched official torch.nn.LayerNorm within 1e-5 precision!\")\n", "else:\n", " print(\"\u26a0\ufe0f PyTorch not installed in this environment; NumPy LayerNorm verified.\")\n" ] }, { "cell_type": "markdown", "id": "2bac529a", "metadata": {}, "source": [ "---\n", "## 7. Multi-Head Self-Attention (Transformer Core)\n", "\n", "Multi-Head Attention (Vaswani et al., 2017) is the central engine of the Transformer. It enables the model to simultaneously attend to information from different representation subspaces at different positions.\n", "\n", "### Mathematical Formulation\n", "Given sequence $X \\in \\mathbb{R}^{B \\times T \\times D_{\\text{model}}}$:\n", "1. **Linear Projections:** Project into Query ($Q$), Key ($K$), and Value ($V$):\n", " $$\n", " Q = X W_Q, \\qquad K = X W_K, \\qquad V = X W_V\n", " $$\n", "2. **Multi-Head Reshape:** Split $D_{\\text{model}}$ into $H$ heads each of dimension $d_k = D / H$:\n", " $$\n", " \\text{Shape: } (B, T, D) \\to (B, H, T, d_k)\n", " $$\n", "3. **Scaled Dot-Product Attention:**\n", " $$\n", " \\text{Attention}(Q, K, V) = \\text{softmax}\\left(\\frac{Q K^T}{\\sqrt{d_k}} + M\\right) V\n", " $$\n", " where $M$ is an optional causal mask ($M_{i,j} = -\\infty$ for $j > i$).\n", "4. **Concatenate & Output Projection:**\n", " Combine the $H$ heads back to $(B, T, D)$ and project with $W_O \\in \\mathbb{R}^{D \\times D}$.\n", "\n", "```\n", "Input (B, T, D)\n", " \u251c\u2500\u2500> W_Q \u2500\u2500> Q (B, H, T, d_k) \u2500\u2510\n", " \u251c\u2500\u2500> W_K \u2500\u2500> K (B, H, T, d_k) \u2500\u2500\u2534\u2500\u2500> (Q @ K^T) / sqrt(d_k) + Mask \u2500\u2500> Softmax \u2500\u2500\u2510\n", " \u2514\u2500\u2500> W_V \u2500\u2500> V (B, H, T, d_k) \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2534\u2500\u2500> Attention (B, H, T, d_k)\n", " \u2502\n", " Reshape & Concat <\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2518\n", " \u2502\n", " W_O \u2500\u2500> Output (B, T, D)\n", "```\n", "\n", "> **\ud83d\udca1 Interview Gotchas & Pro-Tips:**\n", "> - **Why scale by $\\frac{1}{\\sqrt{d_k}}$?** For large $d_k$, the dot product $\\sum_{i=1}^{d_k} q_i k_i$ has variance equal to $d_k$. Without scaling, the values explode in magnitude, pushing softmax into regions with near-zero gradients (vanishing gradient problem).\n", "> - **Causal Mask Implementation:** Use `torch.tril(torch.ones(T, T))` and `.masked_fill(mask == 0, float('-inf'))`.\n", "> - **Efficient Projection:** Instead of creating 3 separate `nn.Linear` layers, top implementations use a single `nn.Linear(D, 3 * D)` and call `.chunk(3, dim=-1)`.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "2adbe6f5", "metadata": {}, "outputs": [], "source": [ "if TORCH_AVAILABLE:\n", " import torch.nn.functional as F\n", "\n", " class MultiHeadSelfAttention(nn.Module):\n", " \"\"\"\n", " Production-quality Multi-Head Attention module from scratch.\n", " Supports causal masking (autoregressive decoding) and batched computation.\n", " \"\"\"\n", " def __init__(self, d_model: int, n_heads: int, dropout: float = 0.0):\n", " super().__init__()\n", " assert d_model % n_heads == 0, \"d_model must be divisible by n_heads\"\n", " \n", " self.d_model = d_model\n", " self.n_heads = n_heads\n", " self.d_k = d_model // n_heads\n", " \n", " # Single combined projection for Q, K, V (3 * d_model) for optimal GPU efficiency\n", " self.qkv_proj = nn.Linear(d_model, 3 * d_model)\n", " self.out_proj = nn.Linear(d_model, d_model)\n", " self.dropout = nn.Dropout(dropout)\n", "\n", " def forward(self, x: torch.Tensor, is_causal: bool = False) -> torch.Tensor:\n", " \"\"\"\n", " x: Input tensor of shape (B, T, D)\n", " is_causal: If True, applies lower-triangular causal attention mask\n", " \"\"\"\n", " B, T, D = x.shape\n", " \n", " # 1. Linear projection -> (B, T, 3 * D)\n", " qkv = self.qkv_proj(x)\n", " \n", " # Split into Q, K, V -> each (B, T, D)\n", " q, k, v = qkv.chunk(3, dim=-1)\n", " \n", " # 2. Reshape and transpose for multi-head: (B, T, H, d_k) -> (B, H, T, d_k)\n", " q = q.view(B, T, self.n_heads, self.d_k).transpose(1, 2)\n", " k = k.view(B, T, self.n_heads, self.d_k).transpose(1, 2)\n", " v = v.view(B, T, self.n_heads, self.d_k).transpose(1, 2)\n", " \n", " # 3. Scaled dot-product: (B, H, T, d_k) @ (B, H, d_k, T) -> (B, H, T, T)\n", " scores = torch.matmul(q, k.transpose(-2, -1)) / (self.d_k ** 0.5)\n", " \n", " # 4. Apply causal mask if required\n", " if is_causal:\n", " causal_mask = torch.tril(torch.ones(T, T, device=x.device, dtype=torch.bool))\n", " scores = scores.masked_fill(~causal_mask, float(\"-inf\"))\n", " \n", " # 5. Softmax along key sequence dimension\n", " attn_weights = F.softmax(scores, dim=-1)\n", " attn_weights = self.dropout(attn_weights)\n", " \n", " # 6. Weighted sum over values: (B, H, T, T) @ (B, H, T, d_k) -> (B, H, T, d_k)\n", " context = torch.matmul(attn_weights, v)\n", " \n", " # 7. Concatenate heads: (B, H, T, d_k) -> (B, T, H, d_k) -> (B, T, D)\n", " context = context.transpose(1, 2).contiguous().view(B, T, D)\n", " \n", " # 8. Final linear projection\n", " return self.out_proj(context)\n" ] }, { "cell_type": "code", "execution_count": null, "id": "fac439c4", "metadata": {}, "outputs": [], "source": [ "# --- Verification: Multi-Head Self-Attention ---\n", "if TORCH_AVAILABLE:\n", " torch.manual_seed(42)\n", " B, T, D, H = 2, 6, 64, 4\n", " x = torch.randn(B, T, D)\n", " \n", " mha = MultiHeadSelfAttention(d_model=D, n_heads=H)\n", " \n", " # Test bidirectional attention\n", " out_bidi = mha(x, is_causal=False)\n", " assert out_bidi.shape == (B, T, D), f\"Expected {(B, T, D)}, got {out_bidi.shape}\"\n", " \n", " # Test causal attention\n", " out_causal = mha(x, is_causal=True)\n", " assert out_causal.shape == (B, T, D), f\"Expected {(B, T, D)}, got {out_causal.shape}\"\n", " \n", " # Test that causal mask prevents future tokens from affecting past tokens\n", " x_perturbed = x.clone()\n", " x_perturbed[:, -1, :] += 10.0\n", " out_perturbed_causal = mha(x_perturbed, is_causal=True)\n", " \n", " token_0_diff = (out_causal[:, 0, :] - out_perturbed_causal[:, 0, :]).abs().max().item()\n", " assert token_0_diff < 1e-6, f\"Causal mask failed! Future token affected position 0: diff={token_0_diff}\"\n", " \n", " print(f\"MHA Output Shape: {out_causal.shape}\")\n", " print(\"\u2705 Multi-Head Self-Attention passed all verification & causality checks!\")\n", "else:\n", " print(\"\u26a0\ufe0f PyTorch not installed; skipped MHA execution.\")\n" ] }, { "cell_type": "markdown", "id": "233347ca", "metadata": {}, "source": [ "---\n", "## 8. Low-Rank Adaptation (LoRA) Layer\n", "\n", "LoRA (Hu et al., 2021) is the standard parameter-efficient fine-tuning (PEFT) method used to adapt large foundation models (LLaMA, Mistral) by training $< 1\\%$ of total parameters.\n", "\n", "### Mathematical Formulation\n", "For a pre-trained linear layer $h = W_0 x$, LoRA decomposes the weight update $\\Delta W$ into two low-rank matrices:\n", "$$\n", "h = W_0 x + \\Delta W x = W_0 x + \\frac{\\alpha}{r} (B A) x\n", "$$\n", "\n", "Where:\n", "- $W_0 \\in \\mathbb{R}^{d_{\\text{out}} \\times d_{\\text{in}}}$ is **frozen** (`requires_grad = False`).\n", "- $A \\in \\mathbb{R}^{r \\times d_{\\text{in}}}$: down-projection with rank $r \\ll \\min(d_{\\text{in}}, d_{\\text{out}})$.\n", "- $B \\in \\mathbb{R}^{d_{\\text{out}} \\times r}$: up-projection.\n", "- $\\alpha$: constant scaling hyperparameter.\n", "- $\\frac{\\alpha}{r}$: scaling factor to keep learning dynamics stable across different rank choices.\n", "\n", "```\n", " Input x (..., d_in)\n", " \u251c\u2500\u2500\u2500> [Frozen W_0] (d_out x d_in) \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2510\n", " \u2502 \u2502 (+) \u2500\u2500> Output (..., d_out)\n", " \u2514\u2500\u2500\u2500> [Trainable A] (r x d_in) \u2500\u2500> [Trainable B] (d_out x r) * (alpha / r) \u2500\u2500\u2518\n", "```\n", "\n", "> **\ud83d\udca1 The Two Critical Interview Traps of LoRA:**\n", "> 1. **The Initialization Trap:**\n", "> - Matrix $A$ is initialized with random Gaussian / Kaiming distribution.\n", "> - Matrix $B$ **MUST be initialized to exact zeros** (`nn.init.zeros_(self.lora_B)`).\n", "> - **Why?** Because at step 0, $\\Delta W = B \\times A = 0 \\times A = 0$. This guarantees that before training begins, the adapted model behaves **identically** to the base model with zero modification!\n", "> 2. **Zero-Inference Overhead (Weight Folding):**\n", "> - During deployment, compute $W_{\\text{merged}} = W_0 + \\frac{\\alpha}{r} B A$.\n", "> - LoRA adds **zero latency** and **zero extra FLOPs** during production inference!\n" ] }, { "cell_type": "code", "execution_count": null, "id": "0da2cd2a", "metadata": {}, "outputs": [], "source": [ "if TORCH_AVAILABLE:\n", " import math\n", "\n", " class LoRALinear(nn.Module):\n", " \"\"\"\n", " Low-Rank Adaptation (LoRA) Linear Layer wrapper from scratch.\n", " \"\"\"\n", " def __init__(self, in_features: int, out_features: int, rank: int = 8, alpha: float = 16.0):\n", " super().__init__()\n", " self.in_features = in_features\n", " self.out_features = out_features\n", " self.rank = rank\n", " self.scaling = alpha / rank\n", " self.merged = False\n", "\n", " # Base frozen weights\n", " self.weight = nn.Parameter(torch.randn(out_features, in_features), requires_grad=False)\n", " self.bias = nn.Parameter(torch.zeros(out_features), requires_grad=False)\n", "\n", " # Low-rank adapter matrices\n", " self.lora_A = nn.Parameter(torch.empty(rank, in_features))\n", " self.lora_B = nn.Parameter(torch.empty(out_features, rank))\n", "\n", " self.reset_parameters()\n", "\n", " def reset_parameters(self):\n", " # 1. Kaiming uniform for A\n", " nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))\n", " # 2. Crucial: B initialized to EXACT ZERO so delta_W starts at 0!\n", " nn.init.zeros_(self.lora_B)\n", "\n", " def forward(self, x: torch.Tensor) -> torch.Tensor:\n", " if self.merged:\n", " return F.linear(x, self.weight, self.bias)\n", " \n", " # Base linear pass: x @ W_0^T + b\n", " base_out = F.linear(x, self.weight, self.bias)\n", " \n", " # LoRA branch: (x @ A^T) @ B^T * scaling\n", " lora_out = (x @ self.lora_A.T) @ self.lora_B.T * self.scaling\n", " return base_out + lora_out\n", "\n", " def merge_weights(self):\n", " \"\"\"Fold LoRA delta into base weight for zero-latency inference.\"\"\"\n", " if not self.merged:\n", " delta_w = (self.lora_B @ self.lora_A) * self.scaling\n", " self.weight.data += delta_w\n", " self.merged = True\n", "\n", " def unmerge_weights(self):\n", " \"\"\"Unfold LoRA weights back to continue fine-tuning.\"\"\"\n", " if self.merged:\n", " delta_w = (self.lora_B @ self.lora_A) * self.scaling\n", " self.weight.data -= delta_w\n", " self.merged = False\n" ] }, { "cell_type": "code", "execution_count": null, "id": "afdfb68d", "metadata": {}, "outputs": [], "source": [ "# --- Verification: LoRA Layer ---\n", "if TORCH_AVAILABLE:\n", " torch.manual_seed(42)\n", " in_dim, out_dim, rank = 128, 64, 4\n", " x_test = torch.randn(8, in_dim)\n", " \n", " lora_layer = LoRALinear(in_features=in_dim, out_features=out_dim, rank=rank, alpha=8.0)\n", " \n", " # 1. Check parameter counts\n", " trainable_params = sum(p.numel() for p in lora_layer.parameters() if p.requires_grad)\n", " frozen_params = sum(p.numel() for p in lora_layer.parameters() if not p.requires_grad)\n", " print(f\"Frozen parameters (W_0, bias): {frozen_params}\")\n", " print(f\"Trainable LoRA parameters (A, B): {trainable_params} ({trainable_params / (frozen_params + trainable_params) * 100:.2f}%)\")\n", " \n", " # 2. Check Initialization Trap: initial output must EQUAL base linear output exactly\n", " base_linear_out = F.linear(x_test, lora_layer.weight, lora_layer.bias)\n", " initial_lora_out = lora_layer(x_test)\n", " assert torch.allclose(initial_lora_out, base_linear_out), \"LoRA initialization trap violated! Output != base output at step 0\"\n", " \n", " # 3. Simulate training update: modify lora_B\n", " with torch.no_grad():\n", " lora_layer.lora_B.add_(torch.randn_like(lora_layer.lora_B) * 0.1)\n", " \n", " out_before_merge = lora_layer(x_test)\n", " \n", " # 4. Test weight merging equivalence\n", " lora_layer.merge_weights()\n", " out_after_merge = lora_layer(x_test)\n", " assert torch.allclose(out_before_merge, out_after_merge, atol=1e-5), \"Merged output does not match unmerged output!\"\n", " \n", " print(\"\u2705 LoRA Layer passed parameter-efficiency, initialization trap, and weight-merging checks!\")\n", "else:\n", " print(\"\u26a0\ufe0f PyTorch not installed; skipped LoRA execution.\")\n" ] }, { "cell_type": "markdown", "id": "bbd60ee4", "metadata": {}, "source": [ "---\n", "## 9. Byte-Pair Encoding (BPE) Tokenizer\n", "\n", "Byte-Pair Encoding (Gage 1994, Sennrich et al. 2015) is the subword tokenization algorithm powering modern LLMs (GPT-2, GPT-3, GPT-4, LLaMA). It provides the optimal trade-off between character-level (too long sequences) and word-level (out-of-vocabulary words).\n", "\n", "### Algorithmic Steps\n", "1. **Base Vocabulary:** Split each word in the training corpus into characters, appending an end-of-word marker ``:\n", " `\"low\" -> ('l', 'o', 'w', '')`\n", "2. **Frequency Counting:** Count all adjacent symbol pairs across the entire dictionary:\n", " `('l', 'o') -> 5, ('o', 'w') -> 5, ('w', '') -> 5`\n", "3. **Merge Step:** Identify the most frequent pair $(c_1, c_2)$ and merge it into a single token $c_1c_2$:\n", " `('l', 'o', 'w', '') -> ('lo', 'w', '')`\n", "4. **Iterate:** Repeat for $K$ merge steps or until target vocabulary size is reached.\n", "5. **Inference (Tokenizing New Text):** Apply the learned merge rules **in the exact chronological order** they were learned.\n", "\n", "```\n", "Initial: ['l', 'o', 'w', ''] ['l', 'o', 'w', 'e', 's', 't', '']\n", "Merge 1: ['lo', 'w', ''] ['lo', 'w', 'e', 's', 't', '']\n", "Merge 2: ['low', ''] ['low', 'e', 's', 't', '']\n", "Merge 3: ['low', ''] ['low', 'est', '']\n", "```\n", "\n", "> **\ud83d\udca1 Interview Gotchas & Pro-Tips:**\n", "> - **End-of-word marker ``:** Crucial for distinguishing between substrings inside words vs at the ends (e.g. \"th\" at the start of \"the\" vs at the end of \"with\").\n", "> - **Subword handling of unseen/OOD words:** If a test word has never been seen before, it gracefully decomposes into individual characters without throwing ``!\n" ] }, { "cell_type": "code", "execution_count": null, "id": "c4948fbd", "metadata": {}, "outputs": [], "source": [ "import collections\n", "import re\n", "from typing import Dict, Tuple, List\n", "\n", "class BPETokenizer:\n", " \"\"\"\n", " Byte-Pair Encoding (BPE) Tokenizer implemented from scratch in pure Python.\n", " \"\"\"\n", " def __init__(self, num_merges: int = 10):\n", " self.num_merges = num_merges\n", " self.merges: List[Tuple[str, str]] = []\n", " self.vocab: List[str] = []\n", "\n", " @staticmethod\n", " def _get_stats(vocab: Dict[Tuple[str, ...], int]) -> Dict[Tuple[str, str], int]:\n", " \"\"\"Count frequency of all adjacent symbol pairs across vocabulary.\"\"\"\n", " pairs = collections.defaultdict(int)\n", " for word_tuple, freq in vocab.items():\n", " for i in range(len(word_tuple) - 1):\n", " pairs[(word_tuple[i], word_tuple[i + 1])] += freq\n", " return pairs\n", "\n", " @staticmethod\n", " def _merge_vocab(pair: Tuple[str, str], vocab: Dict[Tuple[str, ...], int]) -> Dict[Tuple[str, ...], int]:\n", " \"\"\"Merge all occurrences of pair in vocabulary words.\"\"\"\n", " v_out = {}\n", " first, second = pair\n", " new_token = first + second\n", " \n", " for word_tuple, freq in vocab.items():\n", " new_word = []\n", " i = 0\n", " while i < len(word_tuple):\n", " if i < len(word_tuple) - 1 and word_tuple[i] == first and word_tuple[i + 1] == second:\n", " new_word.append(new_token)\n", " i += 2\n", " else:\n", " new_word.append(word_tuple[i])\n", " i += 1\n", " v_out[tuple(new_word)] = freq\n", " return v_out\n", "\n", " def fit(self, corpus: str) -> \"BPETokenizer\":\n", " words = re.findall(r\"\\S+\", corpus)\n", " word_freqs = collections.Counter(words)\n", " \n", " vocab = {tuple(list(word) + [\"\"]): freq for word, freq in word_freqs.items()}\n", " \n", " self.merges = []\n", " for step in range(self.num_merges):\n", " pairs = self._get_stats(vocab)\n", " if not pairs:\n", " break\n", " best_pair = max(pairs, key=pairs.get)\n", " vocab = self._merge_vocab(best_pair, vocab)\n", " self.merges.append(best_pair)\n", " \n", " vocab_set = set()\n", " for word_tuple in vocab.keys():\n", " for token in word_tuple:\n", " vocab_set.add(token)\n", " self.vocab = sorted(list(vocab_set))\n", " return self\n", "\n", " def tokenize(self, text: str) -> List[str]:\n", " \"\"\"Tokenize an unseen input string using the learned merge rules in chronological order.\"\"\"\n", " words = re.findall(r\"\\S+\", text)\n", " tokens = []\n", " \n", " for word in words:\n", " word_tuple = tuple(list(word) + [\"\"])\n", " \n", " for pair in self.merges:\n", " first, second = pair\n", " new_token = first + second\n", " new_word = []\n", " i = 0\n", " while i < len(word_tuple):\n", " if i < len(word_tuple) - 1 and word_tuple[i] == first and word_tuple[i + 1] == second:\n", " new_word.append(new_token)\n", " i += 2\n", " else:\n", " new_word.append(word_tuple[i])\n", " i += 1\n", " word_tuple = tuple(new_word)\n", " \n", " tokens.extend(list(word_tuple))\n", " \n", " return tokens\n" ] }, { "cell_type": "code", "execution_count": null, "id": "2c9d20d6", "metadata": {}, "outputs": [], "source": [ "# --- Verification & Unit Test: BPE Tokenizer ---\n", "corpus = \"low low low low low lower lower newest newest widest widest widest\"\n", "bpe = BPETokenizer(num_merges=6)\n", "bpe.fit(corpus)\n", "\n", "print(\"Top 6 Learned BPE Merges in order:\")\n", "for i, merge in enumerate(bpe.merges, 1):\n", " print(f\" {i}. {merge} -> '{merge[0] + merge[1]}'\")\n", "\n", "test_sentence = \"lower lowest\"\n", "tokens = bpe.tokenize(test_sentence)\n", "print(f\"\\nTokenized '{test_sentence}':\\n {tokens}\")\n", "\n", "assert len(bpe.merges) == 6, \"Expected 6 merges to be recorded\"\n", "assert any(\"low\" in t for t in tokens), \"Expected 'low' subword in tokenization\"\n", "print(\"\u2705 Byte-Pair Encoding passed all verification checks!\")\n" ] }, { "cell_type": "markdown", "id": "7820dcd4", "metadata": {}, "source": [ "---\n", "## 10. MLE Interview Quick Reference & Cheat Sheet\n", "\n", "### Time & Space Complexity Summary\n", "| Algorithm | Train Time | Inference Time | Space Complexity |\n", "|---|---|---|---|\n", "| **Linear Reg (GD)** | $\\mathcal{O}(\\text{iters} \\cdot N D)$ | $\\mathcal{O}(D)$ | $\\mathcal{O}(D)$ |\n", "| **Linear Reg (Normal Eq)** | $\\mathcal{O}(D^3 + N D^2)$ | $\\mathcal{O}(D)$ | $\\mathcal{O}(D^2)$ |\n", "| **Logistic Regression** | $\\mathcal{O}(\\text{iters} \\cdot N D)$ | $\\mathcal{O}(D)$ | $\\mathcal{O}(D)$ |\n", "| **K-Means** | $\\mathcal{O}(\\text{iters} \\cdot N K D)$ | $\\mathcal{O}(K D)$ | $\\mathcal{O}(N D + K D)$ |\n", "| **KNN** | $\\mathcal{O}(1)$ (Lazy) | $\\mathcal{O}(N D + N)$ | $\\mathcal{O}(N D)$ |\n", "| **Multi-Head Attention** | $\\mathcal{O}(T^2 D)$ | $\\mathcal{O}(T^2 D)$ (or $\\mathcal{O}(T D)$ w/ KV Cache) | $\\mathcal{O}(B T D + B H T^2)$ |\n", "| **LayerNorm** | $\\mathcal{O}(N D)$ | $\\mathcal{O}(N D)$ | $\\mathcal{O}(D)$ |\n", "| **LoRA** | $\\mathcal{O}(r (d_{\\text{in}} + d_{\\text{out}}))$ | $\\mathcal{O}(1)$ after fold | $\\mathcal{O}(r (d_{\\text{in}} + d_{\\text{out}}))$ |\n", "\n", "---\n", "\n", "### Top Whiteboard Traps Checklist\n", "1. **Linear Reg:** Check for collinearity $\\to$ recommend `np.linalg.pinv`.\n", "2. **Logistic Reg:** Watch out for `np.log(0)` $\\to$ use `np.clip` and branch sigmoid for $z \\ge 0$ vs $z < 0$.\n", "3. **K-Means:** Never use nested loops for distances $\\to$ use NumPy broadcasting. Initialize via random choice without replacement and update centroids via list comprehension.\n", "4. **KNN:** Don't use `np.argsort` ($\\mathcal{O}(N \\log N)$) $\\to$ use `np.argpartition` ($\\mathcal{O}(N)$).\n", "5. **Backpropagation:** Remember He initialization for ReLU to break symmetry.\n", "6. **LayerNorm:** Normalize over the **feature dimension** $D$, NOT the batch dimension!\n", "7. **Multi-Head Attention:** Scale dot products by $\\frac{1}{\\sqrt{d_k}}$ to avoid vanishing gradients in softmax.\n", "8. **LoRA:** Matrix $B$ **must be initialized to zero** so initial adaptation is zero.\n", "9. **BPE:** Always preserve the order of learned merge rules at inference time.\n" ] } ], "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.9.6" } }, "nbformat": 4, "nbformat_minor": 5 }