workshop/1b_neural_nets.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b1789f5",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "!pip3 install \"d2l==1.0.3\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c3fb25b",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Neural Networks\n",
    "\n",
    "## Multilayer Perceptrons"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe983b8c",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "source": [
    "### Limits of Linear Models\n",
    "\n",
    "<!-- ![OR vs XOR Graphic](../images/xor_problem.png) -->\n",
    "<img src=\"../images/xor_problem.png\" alt=\"OR vs XOR Graphic\" width=\"500\"/>\n",
    "\n",
    "XOR involves two sets of 2D points: \n",
    "  - (0,0), (1,1) belong to one class.\n",
    "  - (0,1), (1,0) belong to the other class.\n",
    "  \n",
    "Linear limitation\n",
    "  - No single linear line can separate these classes.\n",
    "  \n",
    "Key takeaway\n",
    "  - XOR is not linearly separable.\n",
    "  - Nonlinear models (e.g., neural networks) are required.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71b96513",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Multi-layer perceptron"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8049856d",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "source": [
    "<!-- ![An MLP with a hidden layer of five hidden units.](../d2l-en/img/mlp.svg) -->\n",
    "<img src=\"../d2l-en/img/mlp.svg\" alt=\"An MLP with a hidden layer of five hidden units.\" width=\"500\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e0615c5",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "source": [
    "**Notations**:\n",
    "- Input: $\\mathbf{X} \\in \\mathbb{R}^{1 \\times d}$\n",
    "- Hidden: $\\mathbf{H} \\in \\mathbb{R}^{1 \\times h}$\n",
    "- Output: $\\mathbf{O} \\in \\mathbb{R}^{1 \\times q}$\n",
    "\n",
    "**Parameters**:\n",
    "- $\\mathbf{W}^{(1)} \\in \\mathbb{R}^{d \\times h}$, $\\mathbf{b}^{(1)} \\in \\mathbb{R}^{1 \\times h}$\n",
    "- $\\mathbf{W}^{(2)} \\in \\mathbb{R}^{h \\times q}$, $\\mathbf{b}^{(2)} \\in \\mathbb{R}^{1 \\times q}$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97f355a7",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "source": [
    "**Forward Pass**:\n",
    "$$\n",
    "\\mathbf{H} = \\mathbf{X} \\mathbf{W}^{(1)} + \\mathbf{b}^{(1)}\n",
    "$$\n",
    "$$\n",
    "\\mathbf{O} = \\mathbf{H} \\mathbf{W}^{(2)} + \\mathbf{b}^{(2)}\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3c5a684",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "**Affine Function Collapse**\n",
    "- Equivalent single-layer model:\n",
    "$$\n",
    "\\mathbf{O} = \\mathbf{X} \\mathbf{W} + \\mathbf{b}\n",
    "$$\n",
    "with $\\mathbf{W} = \\mathbf{W}^{(1)} \\mathbf{W}^{(2)}$ and $\\mathbf{b} = \\mathbf{b}^{(1)} \\mathbf{W}^{(2)} + \\mathbf{b}^{(2)}$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f5c6d26",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "source": [
    "**Nonlinear Activation**\n",
    "- Introduce activation function $\\sigma(x)$:\n",
    "$$\n",
    "\\mathbf{H} = \\sigma(\\mathbf{X} \\mathbf{W}^{(1)} + \\mathbf{b}^{(1)})\n",
    "$$\n",
    "$$\n",
    "\\mathbf{O} = \\mathbf{H} \\mathbf{W}^{(2)} + \\mathbf{b}^{(2)}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c176030e",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "source": [
    "**Multi-Layer Stacking**\n",
    "$$\n",
    "\\mathbf{H}^{(1)} = \\sigma_1(\\mathbf{X} \\mathbf{W}^{(1)} + \\mathbf{b}^{(1)})\n",
    "$$\n",
    "$$\n",
    "\\mathbf{H}^{(2)} = \\sigma_2(\\mathbf{H}^{(1)} \\mathbf{W}^{(2)} + \\mathbf{b}^{(2)})\n",
    "$$\n",
    "$$\n",
    "\\dots\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21ae1dbf",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Universal Approximation Theorem\n",
    "\n",
    "Let $C(X, \\mathbb{R}^m)$ denote the set of continuous functions from a subset $X$ of Euclidean $\\mathbb{R}^n$ to Euclidean space $\\mathbb{R}^m$.  \n",
    "Let $\\sigma \\in C(\\mathbb{R}, \\mathbb{R})$ be a continuous non-polynomial activation function. \n",
    "\n",
    "For every:\n",
    "- $m, n \\in \\mathbb{N}$\n",
    "- compact $X \\subseteq \\mathbb{R}^n$\n",
    "- $f \\in C(X, \\mathbb{R}^m)$\n",
    "- $\\varepsilon > 0$\n",
    "\n",
    "There exist:\n",
    "- $k \\in \\mathbb{N}$\n",
    "- $A \\in \\mathbb{R}^{k \\times n}$\n",
    "- $b \\in \\mathbb{R}^k$\n",
    "- $C \\in \\mathbb{R}^{m \\times k}$\n",
    "\n",
    "Such that:\n",
    "$$\n",
    "\\sup_{x \\in X} \\|f(x) - g(x)\\| < \\varepsilon\n",
    "$$\n",
    "Where:\n",
    "$$\n",
    "g(x) = C \\cdot \\left(\\sigma \\circ (A \\cdot x + b)\\right)\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c60d0a7",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Activation functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "cec00ab5",
   "metadata": {
    "origin_pos": 3,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import torch\n",
    "from d2l import torch as d2l"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76090679",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "source": [
    "ReLU provides a very simple nonlinear transformation:\n",
    "$$\\operatorname{ReLU}(x) = \\max(x, 0).$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5881b40d",
   "metadata": {
    "origin_pos": 8,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [],
   "source": [
    "x = torch.arange(-8.0, 8.0, 0.1, requires_grad=True)\n",
    "y = torch.relu(x)\n",
    "d2l.plot(x.detach(), y.detach(), 'x', 'relu(x)', figsize=(5, 2.5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d36fda0",
   "metadata": {
    "origin_pos": 13,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [],
   "source": [
    "y.backward(torch.ones_like(x), retain_graph=True)\n",
    "d2l.plot(x.detach(), x.grad, 'x', 'grad of relu', figsize=(5, 2.5))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7df9db18",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "The *sigmoid function* transforms those inputs\n",
    "to outputs that lie on the interval (0, 1):\n",
    "\n",
    "$$\\operatorname{sigmoid}(x) = \\frac{1}{1 + \\exp(-x)}.$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f75a16d6",
   "metadata": {
    "origin_pos": 18,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [],
   "source": [
    "y = torch.sigmoid(x)\n",
    "d2l.plot(x.detach(), y.detach(), 'x', 'sigmoid(x)', figsize=(5, 2.5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4cc9493",
   "metadata": {
    "origin_pos": 23,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [],
   "source": [
    "x.grad.data.zero_()\n",
    "y.backward(torch.ones_like(x),retain_graph=True)\n",
    "d2l.plot(x.detach(), x.grad, 'x', 'grad of sigmoid', figsize=(5, 2.5))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84a05d74",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "The tanh (hyperbolic tangent)\n",
    "function also squashes its inputs\n",
    "between $-1$ and $1$.\n",
    "\n",
    "$$\\operatorname{tanh}(x) = \\frac{1 - \\exp(-2x)}{1 + \\exp(-2x)}.$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59cc5dc6",
   "metadata": {
    "origin_pos": 28,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [],
   "source": [
    "y = torch.tanh(x)\n",
    "d2l.plot(x.detach(), y.detach(), 'x', 'tanh(x)', figsize=(5, 2.5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fef45e56",
   "metadata": {
    "origin_pos": 33,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [],
   "source": [
    "x.grad.data.zero_()\n",
    "y.backward(torch.ones_like(x),retain_graph=True)\n",
    "d2l.plot(x.detach(), x.grad, 'x', 'grad of tanh', figsize=(5, 2.5))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d31d58d9",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Implementation of a Multilayer Perceptron"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "87926c3b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:41:22.340655Z",
     "iopub.status.busy": "2023-08-18T19:41:22.340381Z",
     "iopub.status.idle": "2023-08-18T19:41:25.449640Z",
     "shell.execute_reply": "2023-08-18T19:41:25.448607Z"
    },
    "origin_pos": 3,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from d2l import torch as d2l"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cc10e54",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "source": [
    "Implement an MLP\n",
    "with one hidden layer and 256 hidden units"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "bcccd30d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:41:25.459844Z",
     "iopub.status.busy": "2023-08-18T19:41:25.459149Z",
     "iopub.status.idle": "2023-08-18T19:41:25.472859Z",
     "shell.execute_reply": "2023-08-18T19:41:25.471738Z"
    },
    "origin_pos": 12,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [],
   "source": [
    "class MLPScratch(d2l.Classifier):\n",
    "    def __init__(self, num_inputs, num_outputs, num_hiddens, lr, sigma=0.01):\n",
    "        super().__init__()\n",
    "        self.save_hyperparameters()\n",
    "        self.W1 = nn.Parameter(torch.randn(num_inputs, num_hiddens) * sigma)\n",
    "        self.b1 = nn.Parameter(torch.zeros(num_hiddens))\n",
    "        self.W2 = nn.Parameter(torch.randn(num_hiddens, num_outputs) * sigma)\n",
    "        self.b2 = nn.Parameter(torch.zeros(num_outputs))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5891e5e6",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "Exercise: Implement the ReLU activation\n",
    "\n",
    "Hints:\n",
    "- Create a zero vector using `torch.zeros_like(X)`\n",
    "- Use `torch.max()` with two arguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2af157bb",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:41:25.477923Z",
     "iopub.status.busy": "2023-08-18T19:41:25.477044Z",
     "iopub.status.idle": "2023-08-18T19:41:25.482976Z",
     "shell.execute_reply": "2023-08-18T19:41:25.481963Z"
    },
    "origin_pos": 17,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [],
   "source": [
    "def relu(X):\n",
    "    # TODO\n",
    "    Y = None\n",
    "    return Y\n",
    "\n",
    "relu(torch.Tensor([-2., 0., 1.])) == torch.Tensor([0., 0., 1.])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c679c4c1",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "Exercise: Implement our model\n",
    "\n",
    "$$\n",
    "\\mathbf{H} = \\sigma(\\mathbf{X} \\mathbf{W}^{(1)} + \\mathbf{b}^{(1)})\n",
    "$$\n",
    "$$\n",
    "\\mathbf{O} = \\mathbf{H} \\mathbf{W}^{(2)} + \\mathbf{b}^{(2)}\n",
    "$$\n",
    "\n",
    "Hints:\n",
    "- Use ReLU for $\\sigma$\n",
    "- Use `torch.matmul` or the `@` operator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "7438b40c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:41:25.492513Z",
     "iopub.status.busy": "2023-08-18T19:41:25.491685Z",
     "iopub.status.idle": "2023-08-18T19:41:25.498375Z",
     "shell.execute_reply": "2023-08-18T19:41:25.497344Z"
    },
    "origin_pos": 21,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [],
   "source": [
    "@d2l.add_to_class(MLPScratch)\n",
    "def forward(self, X):\n",
    "    X = X.reshape((-1, self.num_inputs))\n",
    "    # TODO\n",
    "    H = None\n",
    "    Y = None\n",
    "    return Y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18681986",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "The training loop for MLPs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82d57362",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:41:25.502740Z",
     "iopub.status.busy": "2023-08-18T19:41:25.502096Z",
     "iopub.status.idle": "2023-08-18T19:42:19.146140Z",
     "shell.execute_reply": "2023-08-18T19:42:19.144962Z"
    },
    "origin_pos": 23,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [],
   "source": [
    "model = MLPScratch(num_inputs=784, num_outputs=10, num_hiddens=256, lr=0.1)\n",
    "data = d2l.FashionMNIST(batch_size=256)\n",
    "trainer = d2l.Trainer(max_epochs=10)\n",
    "trainer.fit(model, data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e329fd47",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "Concise implementation using `nn.Sequential`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "a5087507",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:42:19.152096Z",
     "iopub.status.busy": "2023-08-18T19:42:19.151778Z",
     "iopub.status.idle": "2023-08-18T19:42:19.157644Z",
     "shell.execute_reply": "2023-08-18T19:42:19.156631Z"
    },
    "origin_pos": 26,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [],
   "source": [
    "class MLP(d2l.Classifier):\n",
    "    def __init__(self, num_outputs, num_hiddens, lr):\n",
    "        super().__init__()\n",
    "        self.save_hyperparameters()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Flatten(),\n",
    "            nn.LazyLinear(num_hiddens),\n",
    "            nn.ReLU(),\n",
    "            nn.LazyLinear(num_outputs),\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4422395",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "source": [
    "The training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "baf391c1",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:42:19.162133Z",
     "iopub.status.busy": "2023-08-18T19:42:19.161837Z",
     "iopub.status.idle": "2023-08-18T19:43:16.475813Z",
     "shell.execute_reply": "2023-08-18T19:43:16.474574Z"
    },
    "origin_pos": 30,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [],
   "source": [
    "model = MLP(num_outputs=10, num_hiddens=256, lr=0.1)\n",
    "trainer.fit(model, data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64f68bd8",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Backpropagation\n",
    "### Forward Propagation: Overview\n",
    "\n",
    "**Input**  \n",
    "  $$\\mathbf{x} \\in \\mathbb{R}^d$$  \n",
    "\n",
    "\n",
    "**Layers**\n",
    "\n",
    "$$\\mathbf{z}= \\mathbf{W}^{(1)} \\mathbf{x}$$\n",
    "$$\\mathbf{h}= \\phi (\\mathbf{z})$$\n",
    "$$\\mathbf{o}= \\mathbf{W}^{(2)} \\mathbf{h}$$\n",
    "\n",
    "**Output**\n",
    "\n",
    "$$\\mathbf{o} = \\mathbf{W}^{(2)} \\phi (\\mathbf{W}^{(1)} \\mathbf{x})$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d5d789e",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Loss and Regularization\n",
    "\n",
    "**Loss Function**:  \n",
    "  $$L = l(\\mathbf{o}, y)$$  \n",
    "\n",
    "**Regularization Term**:  \n",
    "  $$s = \\frac{\\lambda}{2} \\left(\\|\\mathbf{W}^{(1)}\\|_\\textrm{F}^2 + \\|\\mathbf{W}^{(2)}\\|_\\textrm{F}^2\\right)$$  \n",
    "\n",
    "**Objective Function**:  \n",
    "  $$J = L + s$$  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee5e975f",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "source": [
    "### Computational Graph\n",
    "\n",
    "<!-- ![Computational graph of forward propagation.](../d2l-en/img/forward.svg) -->\n",
    "<img src=\"../d2l-en/img/forward.svg\" alt=\"Computational graph of forward propagation.\" width=\"600\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a89a9a56",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Backpropagation: Overview\n",
    "\n",
    "**Weight update by gradient descent**\n",
    "\n",
    "$$\n",
    "w_{ij}^{(k)} \\leftarrow w_{ij}^{(k)} - \\eta \\underbrace{\\frac{\\partial J}{\\partial w_{ij}^{(k)}}}_{\\text{gradient}}\n",
    "$$\n",
    "\n",
    "The step size $\\eta$ is called *learning rate*."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37c8e26e",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "source": [
    "**Gradient Calculation**\n",
    "\n",
    "<!-- Objective Function Derivatives:  \n",
    "$$\\frac{\\partial J}{\\partial L} = 1 \\quad \\textrm{and} \\quad \\frac{\\partial J}{\\partial s} = 1$$ -->\n",
    "\n",
    "Output Layer Gradient:  \n",
    "$$\\frac{\\partial J}{\\partial \\mathbf{o}} = \\frac{\\partial J}{\\partial L} \\cdot \\frac{\\partial L}{\\partial \\mathbf{o}} = \\frac{\\partial L}{\\partial \\mathbf{o}}$$\n",
    "\n",
    "Regularization Term Gradients:  \n",
    "<!-- $$\\frac{\\partial s}{\\partial \\mathbf{W}^{(1)}} = \\lambda \\mathbf{W}^{(1)}, \\quad \\frac{\\partial s}{\\partial \\mathbf{W}^{(2)}} = \\lambda \\mathbf{W}^{(2)}$$ -->\n",
    "$$\\frac{\\partial s}{\\partial w_{ij}^{(1)}} = \\lambda w_{ij}^{(1)}, \\quad \\frac{\\partial s}{\\partial w_{ij}^{(2)}} = \\lambda w_{ij}^{(2)}$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98866315",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "**Backpropagation Through Layers**\n",
    "\n",
    "<!-- Gradient w.r.t. $\\mathbf{W}^{(2)}$:  \n",
    "$$\\frac{\\partial J}{\\partial \\mathbf{W}^{(2)}} = \\frac{\\partial J}{\\partial \\mathbf{o}} \\cdot \\frac{\\partial \\mathbf{o}}{\\partial \\mathbf{W}^{(2)}} + \\frac{\\partial J}{\\partial s} \\cdot \\frac{\\partial s}{\\partial \\mathbf{W}^{(2)}} = \\frac{\\partial L}{\\partial \\mathbf{o}} \\mathbf{h}^\\top + \\lambda \\mathbf{W}^{(2)}$$\n",
    "\n",
    "Hidden Layer Gradient:  \n",
    "$$\\frac{\\partial J}{\\partial \\mathbf{h}} = \\frac{\\partial J}{\\partial \\mathbf{o}} \\cdot \\frac{\\partial \\mathbf{o}}{\\partial \\mathbf{h}} = {\\mathbf{W}^{(2)}}^\\top \\frac{\\partial J}{\\partial \\mathbf{o}}$$\n",
    "\n",
    "Intermediate Variable Gradient:  \n",
    "$$\\frac{\\partial J}{\\partial \\mathbf{z}} = \\frac{\\partial J}{\\partial \\mathbf{h}} \\cdot \\frac{\\partial \\mathbf{h}}{\\partial \\mathbf{z}} = \\frac{\\partial J}{\\partial \\mathbf{h}} \\odot \\phi'(\\mathbf{z})$$\n",
    "\n",
    "Gradient w.r.t. $\\mathbf{W}^{(1)}$:  \n",
    "$$\\frac{\\partial J}{\\partial \\mathbf{W}^{(1)}} = \\frac{\\partial J}{\\partial \\mathbf{z}} \\cdot \\frac{\\partial \\mathbf{z}}{\\partial \\mathbf{W}^{(1)}} + \\frac{\\partial J}{\\partial s} \\cdot \\frac{\\partial s}{\\partial \\mathbf{W}^{(1)}} = \\frac{\\partial J}{\\partial \\mathbf{z}} \\mathbf{x}^\\top + \\lambda \\mathbf{W}^{(1)}$$ -->\n",
    "\n",
    "Gradient w.r.t. $w_{ij}^{(2)}$:  \n",
    "$$\\frac{\\partial J}{\\partial w_{ij}^{(2)}} = \\frac{\\partial J}{\\partial \\mathbf{o}} \\cdot \\frac{\\partial \\mathbf{o}}{\\partial w_{ij}^{(2)}} + \\frac{\\partial J}{\\partial s} \\cdot \\frac{\\partial s}{\\partial w_{ij}^{(2)}} = \\frac{\\partial L}{\\partial o_i} h_j + \\lambda w_{ij}^{(2)}$$\n",
    "\n",
    "Hidden Layer Gradient:  \n",
    "$$\\frac{\\partial J}{\\partial \\mathbf{h}} = \\frac{\\partial J}{\\partial \\mathbf{o}} \\cdot \\frac{\\partial \\mathbf{o}}{\\partial \\mathbf{h}} = \\frac{\\partial L}{\\partial \\mathbf{o}} \\mathbf{W}^{(2)}$$\n",
    "\n",
    "Intermediate Variable Gradient:  \n",
    "$$\\frac{\\partial J}{\\partial \\mathbf{z}} = \\frac{\\partial J}{\\partial \\mathbf{h}} \\cdot \\frac{\\partial \\mathbf{h}}{\\partial \\mathbf{z}} = \\frac{\\partial J}{\\partial \\mathbf{h}} \\operatorname{diag} \\phi'(\\mathbf{z})$$\n",
    "\n",
    "Gradient w.r.t. $w_{ij}^{(1)}$:  \n",
    "$$\\frac{\\partial J}{\\partial w_{ij}^{(1)}} = \\frac{\\partial J}{\\partial \\mathbf{z}} \\cdot \\frac{\\partial \\mathbf{z}}{\\partial w_{ij}^{(1)}} + \\frac{\\partial J}{\\partial s} \\cdot \\frac{\\partial s}{\\partial w_{ij}^{(1)}} = \\frac{\\partial J}{\\partial z_i} x_j + \\lambda w_{ij}^{(1)}$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b68f008",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Training Stability and Initialization\n",
    "\n",
    "- A deep network with $L$ layers: each layer $l$ is defined by a transformation $f_l$ with weights $\\mathbf{W}^{(l)}$.\n",
    "$$\n",
    "\\mathbf{x} \\xrightarrow{f_1} \\mathbf{h}^{(1)} \\xrightarrow{f_2} \\mathbf{h}^{(2)} \\ldots \\xrightarrow{f_L} \\mathbf{h}^{(L)} \\xrightarrow{g} L\n",
    "$$\n",
    "<!-- - Hidden layer output:\n",
    "  $$\\mathbf{h}^{(0)} = \\mathbf{x}$$\n",
    "  $$\\mathbf{h}^{(l)} = f_l (\\mathbf{h}^{(l-1)})$$\n",
    "  $$\\mathbf{o} = \\mathbf{h}^{(L)}$$\n",
    "- so output $\\mathbf{o}$ becomes:\n",
    "  $$\\mathbf{o} = f_L \\circ \\cdots \\circ f_1(\\mathbf{x})$$ -->\n",
    "\n",
    "<!-- - Gradient of $\\mathbf{o}$ w.r.t. $\\mathbf{W}^{(l)}$: -->\n",
    "\n",
    "  <!-- $$\\partial_{\\mathbf{W}^{(l)}} \\mathbf{o} = \\underbrace{\\partial_{\\mathbf{h}^{(L-1)}} \\mathbf{h}^{(L)}}_{ \\mathbf{M}^{(L)} \\stackrel{\\textrm{def}}{=}} \\cdots \\underbrace{\\partial_{\\mathbf{h}^{(l)}} \\mathbf{h}^{(l+1)}}_{ \\mathbf{M}^{(l+1)} \\stackrel{\\textrm{def}}{=}} \\underbrace{\\partial_{\\mathbf{W}^{(l)}} \\mathbf{h}^{(l)}}_{ \\mathbf{v}^{(l)} \\stackrel{\\textrm{def}}{=}}.$$ -->\n",
    "\n",
    "<!-- $$\n",
    "\\frac{\\partial L}{\\partial \\mathbf{h}^{(l)}} = \\frac{\\partial L}{\\partial \\mathbf{h}^{(L)}} \\cdot \\frac{\\partial \\mathbf{h}^{(L)}}{\\partial \\mathbf{h}^{(L-1)}} \\cdots \\frac{\\partial \\mathbf{h}^{(l+1)}}{\\partial \\mathbf{h}^{(l)}}\n",
    "$$ -->\n",
    "\n",
    "- Gradient of $L$ w.r.t. $\\mathbf{W}^{(l)}$:\n",
    "\n",
    "\n",
    "$$\\begin{aligned}\n",
    "\\frac{\\partial L}{\\partial w_{ij}^{(l)}}\n",
    "&= \\frac{\\partial L}{\\partial \\mathbf{h}^{(l)}} \\frac{\\partial \\mathbf{h}^{(l)}}{\\partial w_{ij}^{(l)}} \\\\\n",
    "&= \\frac{\\partial L}{\\partial \\mathbf{h}^{(L)}} \\cdot \\frac{\\partial \\mathbf{h}^{(L)}}{\\partial \\mathbf{h}^{(L-1)}} \\cdots \\frac{\\partial \\mathbf{h}^{(l+1)}}{\\partial \\mathbf{h}^{(l)}} \\frac{\\partial \\mathbf{h}^{(l)}}{\\partial w_{ij}^{(l)}}\n",
    "\\end{aligned}$$\n",
    "\n",
    "\n",
    "- Numerical underflow/overflow due to product of Jacobians with varying determinants.\n",
    "  - **Exploding Gradient**: Large updates, model destruction.\n",
    "  - **Vanishing Gradient**: Small updates, learning becomes impossible."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f27c535b",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Vanishing Gradients\n",
    "\n",
    "- **Sigmoid activation**:  \n",
    "  $$\\sigma(x) = \\frac{1}{1 + \\exp(-x)}$$\n",
    "  \n",
    "- Sigmoid was popular due to its resemblance to biological neurons, but causes gradients to vanish."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc551ba8",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import d2l.torch as d2l\n",
    "\n",
    "x = torch.arange(-8.0, 8.0, 0.1, requires_grad=True)\n",
    "y = torch.sigmoid(x)\n",
    "y.backward(torch.ones_like(x))\n",
    "\n",
    "d2l.plot(x.detach().numpy(), [y.detach().numpy(), x.grad.numpy()],\n",
    "         legend=['sigmoid', 'gradient'], figsize=(4.5, 2.5))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2cf3b10",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Breaking the Symmetry\n",
    "\n",
    "- **Symmetry** in MLPs causes weights to be updated identically, limiting expressiveness.\n",
    "- **Example**: All hidden weights initialized as $\\mathbf{W}^{(1)} = c$ leads to identical outputs and gradients.\n",
    "- **Problem**: Gradient updates maintain symmetry, making hidden units redundant.\n",
    "- **Solution**: Careful initialization or dropout breaks this symmetry.\n",
    "\n",
    "<img src=\"../d2l-en/img/mlp.svg\" alt=\"An MLP with a hidden layer of five hidden units.\" width=\"500\"/>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c050686",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Xavier Initialization\n",
    "\n",
    "- Assume $x_i$ and $w_{ij}$ are normally distributed with zero mean.\n",
    "\n",
    "- For layer output $o_i$:\n",
    "\n",
    "  $$o_{i} = \\sum_{j=1}^{n_\\textrm{in}} w_{ij} x_j$$\n",
    "\n",
    "- Variance of the output:\n",
    "\n",
    "  $$\n",
    "  \\begin{aligned}\n",
    "      \\textrm{Var}[o_i] & = E[o_i^2] - (E[o_i])^2 \\\\\n",
    "          & = \\sum_{j=1}^{n_\\textrm{in}} E[w^2_{ij} x^2_j] - 0 \\\\\n",
    "          & = \\sum_{j=1}^{n_\\textrm{in}} E[w^2_{ij}] E[x^2_j] \\\\\n",
    "          & = n_\\textrm{in} \\sigma^2 E[x^2_j].\n",
    "  \\end{aligned}\n",
    "  $$\n",
    "\n",
    "- To stabilize gradients: $n_\\textrm{in} \\sigma^2 = 1$\n",
    "\n",
    "- Xavier Initialization:\n",
    "\n",
    "  $$\\frac{1}{2} (n_\\textrm{in} + n_\\textrm{out}) \\sigma^2 = 1$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c25813d",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "**Xavier Initialization**\n",
    "\n",
    "- Initialize weights by centered normal distribution with variance $\\sigma$:\n",
    "\n",
    "  $$\\sigma = \\sqrt{\\frac{2}{n_\\textrm{in} + n_\\textrm{out}}}$$\n",
    "\n",
    "- Or initialize by uniform distribution:\n",
    "\n",
    "  $$U\\left(-\\sqrt{\\frac{6}{n_\\textrm{in} + n_\\textrm{out}}}, \\sqrt{\\frac{6}{n_\\textrm{in} + n_\\textrm{out}}}\\right)$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb47d7a4",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "Exercise: Apply Uniform Xavier initialization\n",
    "\n",
    "Hints:\n",
    "- Access linear layers using `self.net[i]` with an index `i`\n",
    "- Access the weight of a linear layer using `linear.weight`\n",
    "- Use `nn.init.xavier_uniform_()` with one argument\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e309fb4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import nn\n",
    "from d2l import torch as d2l\n",
    "\n",
    "class MLP(d2l.Classifier):\n",
    "    def __init__(self, num_inputs, num_outputs, num_hiddens):\n",
    "        super().__init__()\n",
    "        self.save_hyperparameters()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Flatten(),\n",
    "            nn.Linear(num_inputs, num_hiddens),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(num_hiddens, num_outputs),\n",
    "        )\n",
    "        # TODO\n",
    "        nn.init.xavier_uniform_(None)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "830ae0c8",
   "metadata": {},
   "source": [
    "**Conclusion: training stability**\n",
    "\n",
    "- Symmetry can lead to ineffective training.\n",
    "- Xavier initialization ensures balanced variance in forward and backward pass\n",
    "- ReLU activations help reduce vanishing gradients and speed convergence."
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Slideshow",
  "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.8.10"
  },
  "required_libs": [],
  "rise": {
   "auto_select": "none",
   "autolaunch": false,
   "enable_chalkboard": true,
   "overlay": "<div class='my-top-right'><img height=80px src='../d2l-en/static/logo-with-text.png'/></div><div class='my-top-left'></div>",
   "scroll": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}