{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f0b42049",
   "metadata": {},
   "source": [
    "# Lesson 31: Multi-Layer Perceptrons — Composing Projections\n",
    "\n",
    "Lesson 30's single neuron, trained however carefully, could not beat chance on the ring-inside-a-disk dataset — a single linear projection just isn't expressive enough. This lesson adds one more layer: instead of *one* projection followed by a threshold, use *two* projections with a nonlinearity in between. That's it. That's the entire idea behind every deep network in this course — stack more of these."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cca3cfad",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70d06988",
   "metadata": {},
   "source": [
    "## The unsolvable problem, again"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4731bcc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = np.random.default_rng(0)\n",
    "theta_in = rng.uniform(0, 2 * np.pi, 60)\n",
    "inner = np.stack([0.5 * np.cos(theta_in), 0.5 * np.sin(theta_in)], axis=1) + rng.normal(0, 0.1, (60, 2))\n",
    "theta_out = rng.uniform(0, 2 * np.pi, 60)\n",
    "outer = np.stack([2.0 * np.cos(theta_out), 2.0 * np.sin(theta_out)], axis=1) + rng.normal(0, 0.15, (60, 2))\n",
    "X = np.vstack([inner, outer])\n",
    "y = np.concatenate([np.zeros(60), np.ones(60)])\n",
    "\n",
    "plt.scatter(*inner.T, s=15, label='inner (y=0)')\n",
    "plt.scatter(*outer.T, s=15, label='outer (y=1)')\n",
    "plt.legend(fontsize=8)\n",
    "plt.gca().set_aspect('equal')\n",
    "plt.title('No single linear projection separates these (Lessons 29-30)')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00f80cdb",
   "metadata": {},
   "source": [
    "## Two projections, one nonlinearity\n",
    "\n",
    "A **multi-layer perceptron (MLP)** chains a *hidden* projection into a new space, a nonlinear **activation** applied elementwise, and then a final linear projection (exactly Lesson 30's neuron) on top of the *transformed* coordinates:\n",
    "\n",
    "$$h = \\text{ReLU}(W_1 x + b_1), \\qquad z = w_2^\\top h + b_2, \\qquad p = \\sigma(z)$$\n",
    "\n",
    "$\\text{ReLU}(u) = \\max(0, u)$ is the simplest common activation: it's linear everywhere except a single kink at zero. That one kink, applied independently to every unit of $h$, is enough — the hidden layer can bend, fold, and stretch the input space so a *linear* boundary in the transformed space corresponds to a highly nonlinear boundary back in the original coordinates."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cddb578c",
   "metadata": {},
   "source": [
    "## Backprop through two layers\n",
    "\n",
    "The chain rule extends cleanly: propagate the error signal $\\partial L/\\partial z$ from Lesson 30 backward through the output projection to get $\\partial L/\\partial h$, then backward through the ReLU and the hidden projection to get $\\partial L/\\partial W_1, \\partial L/\\partial b_1$. Every step is either \"multiply by a weight matrix's transpose\" or \"multiply elementwise by an activation's derivative\" — this is *all* backpropagation ever is, no matter how many layers are stacked."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9bea134",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sigmoid(z):\n",
    "    return 1 / (1 + np.exp(-z))\n",
    "\n",
    "def relu(z):\n",
    "    return np.maximum(0, z)\n",
    "\n",
    "def relu_deriv(z):\n",
    "    return (z > 0).astype(np.float64)\n",
    "\n",
    "def forward(X, W1, b1, W2, b2):\n",
    "    z1 = X @ W1 + b1\n",
    "    a1 = relu(z1)\n",
    "    z2 = (a1 @ W2 + b2).ravel()\n",
    "    p = sigmoid(z2)\n",
    "    return p, (z1, a1, z2)\n",
    "\n",
    "def backward(X, y, p, cache, W2):\n",
    "    z1, a1, z2 = cache\n",
    "    n = len(y)\n",
    "    grad_z2 = ((p - y) / n).reshape(-1, 1)\n",
    "    grad_W2 = a1.T @ grad_z2\n",
    "    grad_b2 = grad_z2.sum(axis=0)\n",
    "    grad_a1 = grad_z2 @ W2.T\n",
    "    grad_z1 = grad_a1 * relu_deriv(z1)\n",
    "    grad_W1 = X.T @ grad_z1\n",
    "    grad_b1 = grad_z1.sum(axis=0)\n",
    "    return grad_W1, grad_b1, grad_W2, grad_b2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42272be4",
   "metadata": {},
   "source": [
    "### Sanity check against PyTorch autograd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4d3aa7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "H = 4\n",
    "init_rng = np.random.default_rng(8)\n",
    "W1 = init_rng.normal(size=(2, H)) * 0.7\n",
    "b1 = np.zeros(H)\n",
    "W2 = init_rng.normal(size=(H, 1)) * 0.7\n",
    "b2 = np.zeros(1)\n",
    "\n",
    "p, cache = forward(X, W1, b1, W2, b2)\n",
    "grad_W1, grad_b1, grad_W2, grad_b2 = backward(X, y, p, cache, W2)\n",
    "\n",
    "X_t = torch.tensor(X)\n",
    "y_t = torch.tensor(y)\n",
    "W1_t = torch.tensor(W1, requires_grad=True)\n",
    "b1_t = torch.tensor(b1, requires_grad=True)\n",
    "W2_t = torch.tensor(W2, requires_grad=True)\n",
    "b2_t = torch.tensor(b2, requires_grad=True)\n",
    "\n",
    "z1_t = X_t @ W1_t + b1_t\n",
    "a1_t = torch.relu(z1_t)\n",
    "z2_t = (a1_t @ W2_t + b2_t).squeeze(-1)\n",
    "loss_t = torch.nn.functional.binary_cross_entropy_with_logits(z2_t, y_t)\n",
    "loss_t.backward()\n",
    "\n",
    "print(f'W1 max diff: {np.abs(grad_W1 - W1_t.grad.numpy()).max():.2e}')\n",
    "print(f'W2 max diff: {np.abs(grad_W2 - W2_t.grad.numpy()).max():.2e}')\n",
    "print(f'b1 max diff: {np.abs(grad_b1 - b1_t.grad.numpy()).max():.2e}')\n",
    "print(f'b2 max diff: {np.abs(grad_b2 - b2_t.grad.numpy()).max():.2e}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a366eb8",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a728c7ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "lr = 0.1\n",
    "losses = []\n",
    "for epoch in range(3000):\n",
    "    p, cache = forward(X, W1, b1, W2, b2)\n",
    "    eps = 1e-9\n",
    "    losses.append(-np.mean(y * np.log(p + eps) + (1 - y) * np.log(1 - p + eps)))\n",
    "    grad_W1, grad_b1, grad_W2, grad_b2 = backward(X, y, p, cache, W2)\n",
    "    W1 -= lr * grad_W1; b1 -= lr * grad_b1\n",
    "    W2 -= lr * grad_W2; b2 -= lr * grad_b2\n",
    "\n",
    "final_pred = (p > 0.5).astype(float)\n",
    "print(f'final accuracy: {(final_pred == y).mean():.1%}  (single neuron, Lesson 30, managed 50%)')\n",
    "\n",
    "plt.plot(losses)\n",
    "plt.xlabel('epoch'); plt.ylabel('loss')\n",
    "plt.title('Training loss (2-layer MLP)')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d016c9d",
   "metadata": {},
   "source": [
    "## What the hidden layer actually did\n",
    "\n",
    "The output layer is *just* a linear projection (Lesson 30's neuron) — but it acts on $h$, the hidden layer's output, not on the original $x$. If the hidden layer did its job, the *transformed* points should already be much easier to separate with a straight line. Since $h$ lives in $\\mathbb{R}^4$ here, we use PCA (Lesson 6) to visualize it in 2D."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "457bc3cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "_, (_, hidden, _) = forward(X, W1, b1, W2, b2)\n",
    "\n",
    "hidden_centered = hidden - hidden.mean(axis=0)\n",
    "cov = np.cov(hidden_centered.T)\n",
    "eigvals, eigvecs = np.linalg.eigh(cov)\n",
    "top2_directions = eigvecs[:, -2:]\n",
    "hidden_2d = hidden_centered @ top2_directions\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(9, 4))\n",
    "axes[0].scatter(*X[y == 0].T, s=15, label='inner')\n",
    "axes[0].scatter(*X[y == 1].T, s=15, label='outer')\n",
    "axes[0].set_aspect('equal')\n",
    "axes[0].set_title('Original input space')\n",
    "axes[0].legend(fontsize=7)\n",
    "\n",
    "axes[1].scatter(*hidden_2d[y == 0].T, s=15, label='inner')\n",
    "axes[1].scatter(*hidden_2d[y == 1].T, s=15, label='outer')\n",
    "axes[1].set_title('Hidden representation\\n(top 2 PCA directions of a 4D space)')\n",
    "axes[1].legend(fontsize=7)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ce98008",
   "metadata": {},
   "source": [
    "Even this lossy 2D snapshot of the full 4D hidden space shows the two classes pulled apart into something close to linearly separable — a big improvement over the 74% ceiling that was the best *any* line could do in the original space (Lesson 29). The actual output neuron works in the full 4D hidden space, where it achieves 100% accuracy exactly. This is the entire mechanism of deep learning in miniature: each layer *reshapes the space* so that the next layer's job gets easier, until the final layer's job is trivial — a single linear projection."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3696d165",
   "metadata": {},
   "source": [
    "### Exercise\n",
    "\n",
    "1. Retrain with `H = 2` (only 2 hidden units) instead of 4, trying a few different `init_rng` seeds. Can 2 hidden units ever reach 100% accuracy on this dataset, or does capacity this limited cap out lower? (Hint: think about how many straight cuts a ReLU layer with $H$ units can combine.)\n",
    "2. Replace `relu`/`relu_deriv` with `tanh`/its derivative ($1 - \\tanh^2$) throughout, and retrain. Does it still solve the problem? Compare the resulting loss curve's shape to the ReLU version's.\n",
    "3. Remove the nonlinearity entirely (replace `relu(z1)` with just `z1` in `forward`, and `relu_deriv(z1)` with an array of ones in `backward`). Confirm the network can no longer beat Lesson 30's ~50% ceiling, and explain algebraically why a linear hidden layer followed by a linear output layer is still just one big linear projection, no matter how many \"layers\" are stacked."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.x"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}