{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c4ae251e",
   "metadata": {},
   "source": [
    "<link rel=\"stylesheet\" href=\"berkeley.css\">\n",
    "\n",
    "<h1 class=\"cal cal-h1\">Lecture 11: Convergence, Momentum, Adam, and Stochastic Gradient Descent (CS 189, Fall 2026)</h1>\n",
    "\n",
    "This notebook accompanies Lecture 11 and Chapter 7 of Bishop and Bishop, *Deep Learning: Foundations and Concepts*. It continues from the Lecture 10 notebook and covers:\n",
    "\n",
    "1. the convergence of gradient descent on a quadratic error surface,\n",
    "2. momentum,\n",
    "3. adaptive learning rates (AdaGrad, RMSProp, and Adam),\n",
    "4. learning rate schedules and warmup, and\n",
    "5. stochastic and mini-batch gradient descent.\n",
    "\n",
    "Sections used in lecture are marked in their headings. **(Lecture Demo n)** is run during the lecture from the matching \"Demo\" slide.\n",
    "\n",
    "| Notebook section | Lecture 11 slide |\n",
    "|---|---|\n",
    "| Recap: The Quadratic Approximation | Recap slides; the eigenvector plots (\"Eigenvectors of the Hessian\") are Lecture Demo 5 in the Lecture 10 notebook |\n",
    "| (Lecture Demo 1) Convergence of Gradient Descent on a Quadratic | Demo: Oscillating convergence; Issues with Gradient Descent (plots) |\n",
    "| (Lecture Demo 2) Gradient Descent with Momentum | Demo: Momentum |\n",
    "| (Lecture Demo 3) Adaptive Learning Rates | Demo: Adam |\n",
    "| (Lecture Demo 4) Learning Rate Schedules | Demo: Learning Rate Schedules |\n",
    "| (Lecture Demo 5) Stochastic Gradient Descent | Demo: Stochastic Gradient Descent |\n",
    "| (Lecture Demo 6) Mini-Batch Stochastic Gradient Descent | Demo: Mini-batch Gradient Descent; How do I Pick the Batch Size? |\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "128d34c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import plotly.express as px\n",
    "import plotly.graph_objects as go\n",
    "from plotly import figure_factory as ff\n",
    "from plotly.subplots import make_subplots\n",
    "\n",
    "colors = px.colors.qualitative.Plotly\n",
    "px.defaults.width = 800"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e65a047b",
   "metadata": {},
   "source": [
    "<link rel=\"stylesheet\" href=\"berkeley.css\">\n",
    "\n",
    "<h2 class=\"cal cal-h2\">Plotting Code</h2>\n",
    "\n",
    "The functions below produce the visualizations used throughout the notebook. They are not part of the course material, but they must be run before the rest of the notebook.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7c9130d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_plot_grid(figs, rows, cols):\n",
    "    \"\"\"Arrange several Plotly figures in a grid of subplots.\n",
    "\n",
    "    The traces and the 2D axis titles of each figure are copied into the grid.\n",
    "    Layout elements such as shapes (e.g., add_vline) and annotations are not\n",
    "    copied and must be added to the returned figure directly.\n",
    "    \"\"\"\n",
    "    def get_trace_type(fig):\n",
    "        for trace in fig.data:\n",
    "            if trace.type == 'surface':\n",
    "                return 'surface'\n",
    "            if trace.type in ('scatter3d', 'mesh3d', 'cone'):\n",
    "                return 'scene'\n",
    "        return 'xy'\n",
    "    types = [get_trace_type(fig) for fig in figs]\n",
    "    specs = [[{'type': ty} for ty in types[i:i + cols]]\n",
    "             for i in range(0, len(figs), cols)]\n",
    "    fig_grid = make_subplots(rows=rows, cols=cols, specs=specs,\n",
    "                             subplot_titles=[fig.layout.title.text or '' for fig in figs])\n",
    "    for i, (fig, ty) in enumerate(zip(figs, types)):\n",
    "        row, col = i // cols + 1, i % cols + 1\n",
    "        fig_grid.add_traces(fig.data, rows=row, cols=col)\n",
    "        if ty == 'xy':\n",
    "            fig_grid.update_xaxes(title_text=fig.layout.xaxis.title.text, row=row, col=col)\n",
    "            fig_grid.update_yaxes(title_text=fig.layout.yaxis.title.text, row=row, col=col)\n",
    "    return fig_grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52e15118",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fixed colors for the true class and fixed symbols for the predicted class.\n",
    "DIAGNOSIS_COLORS = {'benign': 'blue', 'malignant': 'red'}\n",
    "PREDICTION_SYMBOLS = {'Predicted 1 (benign)': 'cross', 'Predicted 0 (malignant)': 'circle-open'}\n",
    "\n",
    "\n",
    "def plot_lr_predictions(w, cancer_df):\n",
    "    \"\"\"Plot the data, the decision boundary, and the predicted probability surface\n",
    "    of the logistic regression model p(t=1|x) = sigmoid(w^T x).\n",
    "\n",
    "    In this dataset t=1 denotes a benign tumor and t=0 a malignant tumor. Color shows the\n",
    "    true class (benign blue, malignant red). The symbol shows the predicted class:\n",
    "    + for a prediction of 1 (benign) and o for a prediction of 0 (malignant).\n",
    "    \"\"\"\n",
    "    df = cancer_df.copy()\n",
    "    p = logistic_model(w, df[['mean radius', 'mean texture']].values)\n",
    "    df['prediction'] = np.where(p > 0.5, 'Predicted 1 (benign)', 'Predicted 0 (malignant)')\n",
    "    fig = px.scatter(df, x='mean radius', y='mean texture',\n",
    "                     symbol='prediction', color='diagnosis',\n",
    "                     symbol_map=PREDICTION_SYMBOLS, color_discrete_map=DIAGNOSIS_COLORS,\n",
    "                     category_orders={'diagnosis': list(DIAGNOSIS_COLORS),\n",
    "                                      'prediction': list(PREDICTION_SYMBOLS)})\n",
    "    for (i, t) in enumerate(fig.data):\n",
    "        t.legendgroup = str(i)\n",
    "    # Decision boundary: w_1 x_1 + w_2 x_2 = 0\n",
    "    xs = np.linspace(df['mean radius'].min(), df['mean radius'].max(), 100)\n",
    "    decision_boundary = -(w[0] * xs) / w[1]\n",
    "    fig.add_scatter(x=xs, y=decision_boundary, mode='lines',\n",
    "                    name='Decision Boundary', legendgroup='Decision Boundary',\n",
    "                    line=dict(color='black', width=2, dash='dash'))\n",
    "    # Probability surface p(t=1|x)\n",
    "    ys = np.linspace(df['mean texture'].min(), df['mean texture'].max(), 100)\n",
    "    X, Y = np.meshgrid(xs, ys)\n",
    "    Z = logistic_model(w, np.c_[X.ravel(), Y.ravel()]).reshape(X.shape)\n",
    "    # Light red where the model predicts malignant (p near 0), white at the decision\n",
    "    # boundary (p = 0.5), and light blue where it predicts benign (p near 1).\n",
    "    fig.add_contour(x=xs, y=ys, z=Z, zmin=0, zmax=1,\n",
    "                    colorscale=[[0, '#fddbc7'], [0.5, '#f7f7f7'], [1, '#d1e5f0']],\n",
    "                    contours=dict(start=0.1, end=0.9, size=0.1), line=dict(width=0.5, color='gray'),\n",
    "                    name='p(t=1|x)',\n",
    "                    colorbar=dict(title='p(t=1|x)', x=1.05, y=0.3, len=0.75))\n",
    "    fig.update_layout(title=f'w=({w[0]:.2f}, {w[1]:.2f})',\n",
    "                      xaxis_range=[xs.min(), xs.max()], yaxis_range=[ys.min(), ys.max()],\n",
    "                      xaxis_title='Mean Radius (standardized)',\n",
    "                      yaxis_title='Mean Texture (standardized)',\n",
    "                      width=800, height=600)\n",
    "    return fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d93950f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_loss(w1, w2, error, ncontours=50, names=('w1', 'w2')):\n",
    "    \"\"\"Plot the error surface (left) and its contours (right) on a grid of weights.\n",
    "\n",
    "    Args:\n",
    "        w1, w2: Meshgrid arrays of the two weights.\n",
    "        error: Array of error values with the same shape as w1 and w2.\n",
    "        ncontours: Number of contour levels.\n",
    "        names: Axis labels for the two weights.\n",
    "    \"\"\"\n",
    "    surf_fig = go.Figure()\n",
    "    surf_fig.add_surface(z=error, x=w1, y=w2,\n",
    "                         colorscale='Viridis_r', opacity=0.7, showscale=False,\n",
    "                         contours=dict(z=dict(show=True, highlightcolor=\"white\",\n",
    "                                              start=error.min(), end=error.max(),\n",
    "                                              size=(error.max() - error.min()) / ncontours)))\n",
    "    surf_fig.update_layout(title=\"Error Surface\")\n",
    "    contour_fig = go.Figure()\n",
    "    contour_fig.add_contour(x=w1.flatten(), y=w2.flatten(), z=error.flatten(),\n",
    "                            colorscale='Viridis_r', opacity=0.7,\n",
    "                            contours=dict(start=error.min(), end=error.max(),\n",
    "                                          size=(error.max() - error.min()) / ncontours),\n",
    "                            colorbar=dict(x=1.05, y=0.35, len=0.75))\n",
    "    contour_fig.update_layout(title=\"Error Contours\")\n",
    "    fig = make_plot_grid([surf_fig, contour_fig], 1, 2).update_layout(height=800)\n",
    "    fig.update_layout(scene=dict(xaxis_title=names[0], yaxis_title=names[1],\n",
    "                                 zaxis_title='Error', aspectmode='cube'))\n",
    "    fig.update_layout(xaxis_range=[w1.min(), w1.max()], yaxis_range=[w2.min(), w2.max()],\n",
    "                      xaxis_title=names[0], yaxis_title=names[1])\n",
    "    return fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91e8df4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_gradient(w1, w2, error, dw1, dw2, scale=1.0, names=('w1', 'w2')):\n",
    "    \"\"\"Plot the error surface and contours together with the gradient field.\n",
    "\n",
    "    The cones (left) and arrows (right) show the gradient vector at each grid\n",
    "    point. The gradient points in the direction of steepest ascent, so gradient\n",
    "    descent moves in the opposite direction.\n",
    "    \"\"\"\n",
    "    fig = plot_loss(w1, w2, error, names=names)\n",
    "    fig.add_trace(\n",
    "        go.Cone(\n",
    "            x=w1.flatten(), y=w2.flatten(), z=np.zeros_like(error).flatten(),  # ground plane\n",
    "            u=dw1.flatten(), v=dw2.flatten(), w=np.zeros_like(error).flatten(),  # no vertical component\n",
    "            sizeref=2, anchor=\"tail\", showscale=False\n",
    "        ), 1, 1)\n",
    "    quiver = ff.create_quiver(\n",
    "        x=w1.flatten(), y=w2.flatten(), u=dw1.flatten(), v=dw2.flatten(),\n",
    "        line_width=2, line_color=\"white\",\n",
    "        scale=scale, arrow_scale=.2, showlegend=False)\n",
    "    fig.add_traces(quiver.data, rows=1, cols=2)\n",
    "    return fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35df01ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_solution_path(fig, errors, ws):\n",
    "    \"\"\"Overlay an optimization path on a figure created by plot_loss or plot_gradient.\"\"\"\n",
    "    s = np.linspace(0, 1, len(ws))\n",
    "    fig.add_scatter3d(x=ws[:, 0], y=ws[:, 1], z=errors, marker_color=s, marker_size=5,\n",
    "                      mode='lines+markers', line=dict(color='black', width=2), opacity=0.5,\n",
    "                      name='Optimization Path', legendgroup='Optimization Path',\n",
    "                      row=1, col=1)\n",
    "    fig.add_scatter(x=ws[:, 0], y=ws[:, 1], marker_color=s,\n",
    "                    mode='lines+markers', line=dict(color='black', width=2), opacity=0.5,\n",
    "                    name='Optimization Path', legendgroup='Optimization Path',\n",
    "                    showlegend=False,\n",
    "                    row=1, col=2)\n",
    "    fig.add_scatter3d(x=[ws[-1, 0]], y=[ws[-1, 1]], z=[errors[-1]],\n",
    "                      mode='markers', marker=dict(color='red', size=10),\n",
    "                      name='Final Solution', legendgroup='Final Solution',\n",
    "                      row=1, col=1)\n",
    "    fig.add_scatter(x=[ws[-1, 0]], y=[ws[-1, 1]],\n",
    "                    mode='markers', marker=dict(color='red', size=20),\n",
    "                    name='Final Solution', legendgroup='Final Solution',\n",
    "                    showlegend=False,\n",
    "                    row=1, col=2)\n",
    "    return fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3207887b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_loss_curve(path, error_func):\n",
    "    \"\"\"Plot the error as a function of the iteration number for a path of weights.\"\"\"\n",
    "    errors = [error_func(w) for w in path]\n",
    "    fig = px.line(x=np.arange(len(errors)), y=errors,\n",
    "                  labels={'x': 'Iteration (gradient steps)', 'y': 'Error'})\n",
    "    fig.update_traces(line_width=4)\n",
    "    fig.update_layout(margin=dict(l=20, r=20, t=20, b=20))\n",
    "    return fig"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de32a25a",
   "metadata": {},
   "source": [
    "<link rel=\"stylesheet\" href=\"berkeley.css\">\n",
    "\n",
    "<h2 class=\"cal cal-h2\">Setup: Models, Error Functions, and Gradients</h2>\n",
    "\n",
    "This section recreates the two problems from the Lecture 10 notebook: logistic regression on two standardized features of the breast cancer data, and non-linear regression with the model $y(x, w) = \\sin(w_0 + w_1 x)$. The derivations of the error functions and gradients are given in that notebook. Each gradient function accepts optional data arguments, so the same function can be used for batch, stochastic, and mini-batch gradient descent.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1fa00fd4",
   "metadata": {},
   "source": [
    "<link rel=\"stylesheet\" href=\"berkeley.css\">\n",
    "\n",
    "<h3 class=\"cal cal-h3\">Logistic Regression</h3>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ce608aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn import datasets\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "\n",
    "cancer_dict = datasets.load_breast_cancer(as_frame=True)\n",
    "cancer_df = pd.DataFrame(cancer_dict.data, columns=cancer_dict.feature_names)\n",
    "cancer_df = cancer_df[['mean radius', 'mean texture']].copy()\n",
    "# Target t=1 denotes a benign tumor and t=0 a malignant tumor.\n",
    "cancer_df['target'] = cancer_dict.target.values\n",
    "cancer_df['diagnosis'] = cancer_dict.target_names[cancer_df['target']]\n",
    "# Standardize each feature to zero mean and unit variance.\n",
    "scaler = StandardScaler()\n",
    "cancer_df[['mean radius', 'mean texture']] = scaler.fit_transform(\n",
    "    cancer_df[['mean radius', 'mean texture']])\n",
    "print(\"The dataset:\", cancer_df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a270f3ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "cancer = dict()\n",
    "cancer['x'] = cancer_df[['mean radius', 'mean texture']].values\n",
    "cancer['t'] = cancer_df['target'].values.astype(float)\n",
    "cancer['npts'] = 30\n",
    "cancer['w1'], cancer['w2'] = np.meshgrid(\n",
    "    np.linspace(-10, 1, cancer['npts']),\n",
    "    np.linspace(-5, 1.3, cancer['npts']))\n",
    "cancer['ws'] = np.stack([cancer['w1'].flatten(), cancer['w2'].flatten()]).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3706232f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sigmoid(z):\n",
    "    \"\"\"Logistic sigmoid function.\"\"\"\n",
    "    return 1 / (1 + np.exp(-z))\n",
    "\n",
    "def logistic_model(w, x):\n",
    "    \"\"\"Logistic regression model p(t=1|x) = sigmoid(w^T x) (no bias term).\"\"\"\n",
    "    return sigmoid(x @ w)\n",
    "\n",
    "\n",
    "def neg_log_likelihood(w, x=None, t=None):\n",
    "    \"\"\"Average negative log-likelihood (cross-entropy) of the logistic model.\n",
    "\n",
    "    Uses the identity -[t ln sigmoid(z) + (1-t) ln(1-sigmoid(z))] = ln(1+e^z) - t z,\n",
    "    evaluated with the numerically stable np.logaddexp(0, z) = ln(1+e^z).\n",
    "    Defaults to the full training set.\n",
    "    \"\"\"\n",
    "    x = cancer['x'] if x is None else x\n",
    "    t = cancer['t'] if t is None else t\n",
    "    z = x @ w\n",
    "    return np.mean(np.logaddexp(0, z) - t * z)\n",
    "\n",
    "\n",
    "def grad_NLL(w, x=None, t=None):\n",
    "    \"\"\"Gradient of the average negative log-likelihood of the logistic model.\n",
    "\n",
    "    Accepts the full data set (default), a mini-batch, or a single data point.\n",
    "    \"\"\"\n",
    "    x = cancer['x'] if x is None else np.atleast_2d(x)\n",
    "    t = cancer['t'] if t is None else np.atleast_1d(t)\n",
    "    y = logistic_model(w, x)\n",
    "    return np.mean((y - t).reshape(-1, 1) * x, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8583cf0",
   "metadata": {},
   "outputs": [],
   "source": [
    "cancer['error'] = np.array([neg_log_likelihood(w) for w in cancer['ws']])\n",
    "cancer['error'] = cancer['error'].reshape(cancer['w1'].shape)\n",
    "best_ind = np.argmin(cancer['error'])\n",
    "cancer['grid_best'] = cancer['ws'][best_ind]\n",
    "\n",
    "(cancer['dw1'], cancer['dw2']) = np.array([grad_NLL(w) for w in cancer['ws']]).T\n",
    "cancer['dw1'] = cancer['dw1'].reshape(cancer['w1'].shape)\n",
    "cancer['dw2'] = cancer['dw2'].reshape(cancer['w1'].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad2c7d0e",
   "metadata": {},
   "source": [
    "<link rel=\"stylesheet\" href=\"berkeley.css\">\n",
    "\n",
    "<h3 class=\"cal cal-h3\">Sine Regression</h3>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8174235",
   "metadata": {},
   "outputs": [],
   "source": [
    "sine = dict()\n",
    "np.random.seed(42)\n",
    "sine['n'] = 200\n",
    "sine['x'] = np.sort(np.random.rand(sine['n']) * 2.5 * np.pi)  # sorted for plotting\n",
    "sine['y'] = np.sin(1.1 + 2.5 * sine['x']) + 0.5 * np.random.randn(sine['n'])\n",
    "sine_df = pd.DataFrame({'x': sine['x'], 'y': sine['y']})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1db88988",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sine_model(w, x):\n",
    "    \"\"\"Non-linear regression model y(x, w) = sin(w_0 + w_1 x).\"\"\"\n",
    "    return np.sin(w[0] + x * w[1])\n",
    "\n",
    "\n",
    "def sine_MSE(w, x=None, y=None):\n",
    "    \"\"\"Mean squared error of the sine model. Defaults to the full training set.\"\"\"\n",
    "    x = sine['x'] if x is None else x\n",
    "    y = sine['y'] if y is None else y\n",
    "    return np.mean((y - sine_model(w, x)) ** 2)\n",
    "\n",
    "\n",
    "def grad_sine_MSE(w, x=None, y=None):\n",
    "    \"\"\"Gradient of the mean squared error of the sine model.\n",
    "\n",
    "    Accepts the full data set (default), a mini-batch, or a single data point.\n",
    "    \"\"\"\n",
    "    x = sine['x'] if x is None else np.atleast_1d(x)\n",
    "    y = sine['y'] if y is None else np.atleast_1d(y)\n",
    "    u = w[0] + w[1] * x\n",
    "    r = y - np.sin(u)\n",
    "    grad_w0 = -2 * np.mean(r * np.cos(u))\n",
    "    grad_w1 = -2 * np.mean(r * np.cos(u) * x)\n",
    "    return np.array([grad_w0, grad_w1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3715c63",
   "metadata": {},
   "outputs": [],
   "source": [
    "sine['npts'] = 30\n",
    "sine['w0'], sine['w1'] = np.meshgrid(\n",
    "    np.linspace(-1.5, 3, sine['npts']), np.linspace(1, 4, sine['npts']))\n",
    "# Stack the grid into an array of weight vectors (w_0, w_1).\n",
    "sine['ws'] = np.stack([sine['w0'].flatten(), sine['w1'].flatten()]).T\n",
    "sine['error'] = np.array([sine_MSE(w) for w in sine['ws']]).reshape(sine['w0'].shape)\n",
    "\n",
    "names = ('w0', 'w1')\n",
    "ind = np.argmin(sine['error'])\n",
    "sine['grid_best'] = sine['ws'][ind, :]\n",
    "\n",
    "(sine['dw0'], sine['dw1']) = np.array([grad_sine_MSE(w) for w in sine['ws']]).T\n",
    "sine['dw0'] = sine['dw0'].reshape(sine['w1'].shape)\n",
    "sine['dw1'] = sine['dw1'].reshape(sine['w1'].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae8bfd92",
   "metadata": {},
   "source": [
    "<link rel=\"stylesheet\" href=\"berkeley.css\">\n",
    "\n",
    "<h3 class=\"cal cal-h3\">Batch Gradient Descent</h3>\n",
    "\n",
    "Batch gradient descent with a fixed learning rate $\\eta$ (Bishop, Eq. 7.16):\n",
    "$$\n",
    "w^{(\\tau)} = w^{(\\tau-1)} - \\eta \\nabla E\\big(w^{(\\tau-1)}\\big).\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0d8f7ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "def gradient_descent(w_0, gradient, learning_rate=1, nepochs=10, epsilon=1e-6):\n",
    "    \"\"\"Batch gradient descent with a fixed learning rate.\n",
    "\n",
    "    Args:\n",
    "        w_0: Initial weight vector.\n",
    "        gradient: Function that returns the gradient of the error at w.\n",
    "        learning_rate: Learning rate (step size) eta.\n",
    "        nepochs: Maximum number of iterations.\n",
    "        epsilon: Stop when the norm of the update falls below this threshold.\n",
    "    Returns:\n",
    "        path: Array containing the weight vector at each iteration.\n",
    "    \"\"\"\n",
    "    w_old = np.asarray(w_0, dtype=float)\n",
    "    path = [w_old]\n",
    "    for _ in range(nepochs):\n",
    "        w = w_old - learning_rate * gradient(w_old)\n",
    "        path.append(w)\n",
    "        if np.linalg.norm(w - w_old) < epsilon:\n",
    "            break\n",
    "        w_old = w\n",
    "    return np.array(path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2880972e",
   "metadata": {},
   "source": [
    "As a reminder, gradient descent on the logistic regression error with $\\eta = 10$, starting from $w^{(0)} = (-10, -5)$:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55b0725b",
   "metadata": {},
   "outputs": [],
   "source": [
    "w0 = np.array([-10., -5.])\n",
    "path = gradient_descent(w0, grad_NLL, learning_rate=10, nepochs=100)\n",
    "\n",
    "fig_loss = make_loss_curve(path, neg_log_likelihood)\n",
    "fig_loss.show()\n",
    "\n",
    "errors = [neg_log_likelihood(w) for w in path]\n",
    "fig = plot_gradient(cancer['w1'], cancer['w2'], cancer['error'],\n",
    "                    cancer['dw1'], cancer['dw2'], scale=2)\n",
    "add_solution_path(fig, errors, path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b91bdefe",
   "metadata": {},
   "source": [
    "<link rel=\"stylesheet\" href=\"berkeley.css\">\n",
    "\n",
    "<h2 class=\"cal cal-h2\">Recap: The Quadratic Approximation</h2>\n",
    "\n",
    "The Hessian, the quadratic approximation of the error function, and the eigenvectors of the Hessian at a stationary point are developed in the Lecture 10 notebook. That notebook also shows that gradient descent on the parabola $f(x) = x^2 + 1$ converges only if $0 < \\eta < 1 = 2/f''(x)$; this notebook generalizes that condition to error functions of many variables. We use the quadratic approximation: near a stationary point $w^\\star$ with Hessian eigenpairs $(\\lambda_i, u_i)$,\n",
    "$$\n",
    "E(w) \\simeq E(w^\\star) + \\frac{1}{2} \\sum_i \\lambda_i \\alpha_i^2, \\qquad \\alpha_i = u_i^\\top (w - w^\\star).\n",
    "$$\n",
    "The eigenvector plots on the slides \"Eigenvectors of the Hessian\" are produced by Lecture Demo 5 (\"The Quadratic Approximation\") in the Lecture 10 notebook.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c5b6e4b",
   "metadata": {},
   "source": [
    "<link rel=\"stylesheet\" href=\"berkeley.css\">\n",
    "\n",
    "<h2 class=\"cal cal-h2\">(Lecture Demo 1) Convergence of Gradient Descent on a Quadratic</h2>\n",
    "\n",
    "Consider gradient descent on the quadratic approximation. In the coordinates $\\alpha_i$ the gradient is $\\partial E / \\partial \\alpha_i = \\lambda_i \\alpha_i$, so each step gives (Bishop, Eqs. 7.26 to 7.29)\n",
    "$$\n",
    "\\alpha_i^{(\\tau)} = (1 - \\eta \\lambda_i)\\, \\alpha_i^{(\\tau-1)} = (1 - \\eta \\lambda_i)^\\tau \\alpha_i^{(0)}.\n",
    "$$\n",
    "The components evolve independently. The iterates converge to $w^\\star$ if and only if $|1 - \\eta \\lambda_i| < 1$ for every $i$, which requires all $\\lambda_i > 0$ and\n",
    "$$\n",
    "0 < \\eta < \\frac{2}{\\lambda_{\\max}}.\n",
    "$$\n",
    "When $1 - \\eta\\lambda_i < 0$ the component $\\alpha_i$ changes sign at every step, which produces oscillation across the valley.\n",
    "\n",
    "For every stable learning rate $\\eta < 2/\\lambda_{\\max}$, the factor of the component along $u_{\\min}$ satisfies $|1 - \\eta\\lambda_{\\min}| > 1 - 2\\lambda_{\\min}/\\lambda_{\\max}$ (Bishop, Eq. 7.30). When the **condition number** $\\kappa = \\lambda_{\\max}/\\lambda_{\\min}$ is large, this bound is close to 1, so that component converges slowly whatever $\\eta$ we choose. The $u_{\\min}$ component is not always the slowest one: for $\\eta$ close to $2/\\lambda_{\\max}$ the factor $|1 - \\eta\\lambda_{\\max}|$ is also close to 1. The step size that minimizes the largest factor is $\\eta^\\star = 2/(\\lambda_{\\max} + \\lambda_{\\min})$, which gives the rate $(\\kappa - 1)/(\\kappa + 1)$ in every direction.\n",
    "\n",
    "*Lecture:* the runs below give the plots on \"Issues with Gradient Descent (GD)\" ($\\eta = 0.1$ and $\\eta = 0.9$) and on \"Demo: Oscillating convergence\".\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94400634",
   "metadata": {},
   "source": [
    "We study the poorly conditioned quadratic\n",
    "$$\n",
    "E(w) = \\frac{1}{2} w^\\top H w, \\qquad H = \\begin{bmatrix} 0.2 & 0 \\\\ 0 & 2 \\end{bmatrix},\n",
    "$$\n",
    "with minimum at $w^\\star = 0$. The eigenvectors are the coordinate axes, $\\lambda_{\\min} = 0.2$, $\\lambda_{\\max} = 2$, and $\\kappa = 10$. Gradient descent converges only for $\\eta < 2/\\lambda_{\\max} = 1$.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8c45c41",
   "metadata": {},
   "outputs": [],
   "source": [
    "H_quad = np.array([[0.2, 0.], [0., 2.]])\n",
    "\n",
    "def quad(w):\n",
    "    \"\"\"Poorly conditioned quadratic E(w) = 0.5 w^T H w (w may be a batch of rows).\"\"\"\n",
    "    return 0.5 * np.sum((w @ H_quad) * w, axis=-1)\n",
    "\n",
    "def grad_quad(w):\n",
    "    \"\"\"Gradient of the quadratic, H w.\"\"\"\n",
    "    return w @ H_quad\n",
    "\n",
    "w1, w2 = np.meshgrid(np.linspace(-5, 5, 30), np.linspace(-5, 5, 30))\n",
    "ws = np.hstack([w1.reshape(-1, 1), w2.reshape(-1, 1)])\n",
    "error = quad(ws).reshape(w1.shape)\n",
    "contour = go.Contour(x=w1.flatten(), y=w2.flatten(), z=error.flatten(), colorscale='Viridis_r',\n",
    "                     contours=dict(start=0, end=20, size=.5))\n",
    "go.Figure(data=contour)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8425377d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_quad_path(path, title=None):\n",
    "    \"\"\"Plot an optimization path on the contours of the quadratic.\"\"\"\n",
    "    fig = go.Figure()\n",
    "    fig.add_trace(contour)\n",
    "    fig.add_scatter(x=path[:, 0], y=path[:, 1],\n",
    "                    mode='lines+markers', line=dict(color='black', width=2),\n",
    "                    marker=dict(size=10, color='black', symbol=\"arrow-bar-up\", angleref=\"previous\"),\n",
    "                    name='Optimization Path', showlegend=False)\n",
    "    fig.update_layout(margin=dict(l=5, r=5, t=40 if title else 5, b=5), title=title)\n",
    "    return fig\n",
    "\n",
    "def report_factors(eta):\n",
    "    \"\"\"Print the per-step contraction factor 1 - eta * lambda_i for each eigen-direction.\"\"\"\n",
    "    for lam_i in np.diag(H_quad):\n",
    "        print(f\"lambda = {lam_i:4.1f}:  1 - eta*lambda = {1 - eta * lam_i:6.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c475bb7",
   "metadata": {},
   "source": [
    "**Small learning rate.** Starting at $(-4, 0)$, only the low-curvature component is non-zero. With $\\eta = 0.1$ its factor is $1 - 0.1 \\times 0.2 = 0.98$, so progress along the valley is slow.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d80dcaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "eta = 0.1\n",
    "report_factors(eta)\n",
    "path = gradient_descent(np.array([-4., 0.]), grad_quad, learning_rate=eta, nepochs=50)\n",
    "plot_quad_path(path, f\"eta = {eta}, start (-4, 0)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5fd5a3d",
   "metadata": {},
   "source": [
    "**Learning rate at the stability limit.** Starting at $(-4, -2)$ with $\\eta = 1 = 2/\\lambda_{\\max}$, the factor along $w_2$ is exactly $-1$. The $w_2$ component alternates between $-2$ and $+2$ indefinitely while the $w_1$ component converges with factor $0.8$.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b139510",
   "metadata": {},
   "outputs": [],
   "source": [
    "eta = 1.0\n",
    "report_factors(eta)\n",
    "path = gradient_descent(np.array([-4., -2.]), grad_quad, learning_rate=eta, nepochs=50)\n",
    "plot_quad_path(path, f\"eta = {eta}, start (-4, -2)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c2cffdf",
   "metadata": {},
   "source": [
    "**Slightly smaller learning rate.** With $\\eta = 0.9$ the factor along $w_2$ is $-0.8$, which gives a damped oscillation across the valley, and the factor along $w_1$ is $0.82$.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57f60d69",
   "metadata": {},
   "outputs": [],
   "source": [
    "eta = 0.9\n",
    "report_factors(eta)\n",
    "path = gradient_descent(np.array([-4., -2.]), grad_quad, learning_rate=eta, nepochs=50)\n",
    "plot_quad_path(path, f\"eta = {eta}, start (-4, -2)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d49aa0bd",
   "metadata": {},
   "source": [
    "**Optimal fixed learning rate.** With $\\eta^\\star = 2/(\\lambda_{\\max} + \\lambda_{\\min}) = 2/2.2 \\approx 0.91$, both factors have magnitude $(\\kappa - 1)/(\\kappa + 1) = 9/11 \\approx 0.82$. No fixed learning rate does better on this problem, and the rate approaches 1 as $\\kappa$ grows. This limitation motivates momentum and adaptive learning rates.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a7774f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "eta = 2 / (H_quad.max() + H_quad[H_quad > 0].min())\n",
    "report_factors(eta)\n",
    "path = gradient_descent(np.array([-4., -2.]), grad_quad, learning_rate=eta, nepochs=50)\n",
    "plot_quad_path(path, f\"eta = {eta:.3f}, start (-4, -2)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "263427dd",
   "metadata": {},
   "source": [
    "<link rel=\"stylesheet\" href=\"berkeley.css\">\n",
    "\n",
    "<h2 class=\"cal cal-h2\">(Lecture Demo 2) Gradient Descent with Momentum</h2>\n",
    "\n",
    "Momentum adds a fraction $\\mu$ of the previous update to the current one (Bishop, Eq. 7.31):\n",
    "$$\n",
    "\\Delta w^{(\\tau-1)} = -\\eta \\nabla E\\big(w^{(\\tau-1)}\\big) + \\mu \\Delta w^{(\\tau-2)},\n",
    "\\qquad\n",
    "w^{(\\tau)} = w^{(\\tau-1)} + \\Delta w^{(\\tau-1)},\n",
    "$$\n",
    "with $0 \\le \\mu < 1$. Substituting for $\\Delta w^{(\\tau-2)}$ repeatedly, with $\\Delta w = 0$ before the first step, gives\n",
    "$$\n",
    "\\Delta w^{(\\tau-1)} = -\\eta \\sum_{k=0}^{\\tau-1} \\mu^k \\, \\nabla E\\big(w^{(\\tau-1-k)}\\big),\n",
    "$$\n",
    "so each step is an exponentially weighted sum of past gradients, and a gradient from $k$ steps ago has weight $\\mu^k$. If the gradient is approximately constant (low curvature), the sum is approximately $\\nabla E/(1 - \\mu)$ and the effective learning rate increases to $\\eta/(1 - \\mu)$. If the gradient changes sign at every step (high curvature), the sum is approximately $\\nabla E/(1 + \\mu)$ and the effective learning rate decreases to $\\eta/(1 + \\mu)$.\n",
    "\n",
    "Deep learning libraries such as PyTorch maintain a **velocity** $v^{(\\tau)} = \\mu v^{(\\tau-1)} + \\nabla E(w^{(\\tau-1)})$ and update $w^{(\\tau)} = w^{(\\tau-1)} - \\eta v^{(\\tau)}$. With $\\Delta w^{(\\tau-1)} = -\\eta v^{(\\tau)}$ this is the same algorithm when $\\eta$ is fixed. The implementation below uses Bishop's form and stores $\\Delta w$ in the variable `dw`.\n",
    "\n",
    "*Lecture:* run from the slide \"Demo: Momentum\".\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "007b8f2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def gd_momentum(w_0, gradient, learning_rate=1, nepochs=10, epsilon=1e-6, momentum=0.9):\n",
    "    \"\"\"Gradient descent with momentum, Bishop Eq. 7.31.\n",
    "\n",
    "    Args:\n",
    "        w_0: Initial weight vector.\n",
    "        gradient: Function that returns the gradient of the error at w.\n",
    "        learning_rate: Learning rate eta.\n",
    "        nepochs: Maximum number of iterations.\n",
    "        epsilon: Stop when the norm of the update falls below this threshold.\n",
    "        momentum: Momentum parameter mu, with 0 <= mu < 1.\n",
    "    Returns:\n",
    "        path: Array containing the weight vector at each iteration.\n",
    "    \"\"\"\n",
    "    w_old = np.asarray(w_0, dtype=float)\n",
    "    path = [w_old]\n",
    "    dw = np.zeros_like(w_old)  # previous update Delta w\n",
    "    for _ in range(nepochs):\n",
    "        dw = momentum * dw - learning_rate * gradient(w_old)\n",
    "        w = w_old + dw\n",
    "        path.append(w)\n",
    "        if np.linalg.norm(w - w_old) < epsilon:\n",
    "            break\n",
    "        w_old = w\n",
    "    return np.array(path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa1d059a",
   "metadata": {},
   "source": [
    "A quick check that the velocity form used by PyTorch gives the same iterates as Bishop's form for a fixed learning rate:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c13a332",
   "metadata": {},
   "outputs": [],
   "source": [
    "def gd_velocity(w_0, gradient, learning_rate=1, nepochs=10, momentum=0.9):\n",
    "    \"\"\"Momentum in the velocity form: v <- mu v + grad, w <- w - eta v.\"\"\"\n",
    "    w = np.asarray(w_0, dtype=float)\n",
    "    path, v = [w], np.zeros_like(w)\n",
    "    for _ in range(nepochs):\n",
    "        v = momentum * v + gradient(w)\n",
    "        w = w - learning_rate * v\n",
    "        path.append(w)\n",
    "    return np.array(path)\n",
    "\n",
    "p1 = gd_momentum(np.array([-4., -2.]), grad_quad, learning_rate=0.5, momentum=0.8, nepochs=20, epsilon=0)\n",
    "p2 = gd_velocity(np.array([-4., -2.]), grad_quad, learning_rate=0.5, momentum=0.8, nepochs=20)\n",
    "print(\"max difference between the two forms:\", np.abs(p1 - p2).max())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd0a6e85",
   "metadata": {},
   "source": [
    "We apply momentum to the toy quadratic function from the previous section, starting at $(-4, -2)$ with $\\eta = 1$ and $\\mu = 0.4$. Without momentum this learning rate lies exactly on the stability limit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0baf74f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "eta, mu = 1.0, 0.4\n",
    "path = gd_momentum(np.array([-4., -2.]), grad_quad, learning_rate=eta, momentum=mu, nepochs=50)\n",
    "plot_quad_path(path, f\"Momentum: eta = {eta}, mu = {mu}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a55221c5",
   "metadata": {},
   "source": [
    "Momentum on the logistic regression error, starting from $w^{(0)} = (0, 0)$. Along the long, flat valley of this error surface the gradient points in a consistent direction, so momentum accelerates the iterates; they overshoot the minimum and then return. Set `momentum=0` to compare with plain gradient descent at the same learning rate.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edca3d51",
   "metadata": {},
   "outputs": [],
   "source": [
    "w0 = np.array([0., 0.])\n",
    "path = gd_momentum(w0, grad_NLL, learning_rate=1., momentum=0.9, nepochs=100)\n",
    "\n",
    "fig_loss = make_loss_curve(path, neg_log_likelihood)\n",
    "fig_loss.show()\n",
    "\n",
    "errors = [neg_log_likelihood(w) for w in path]\n",
    "fig = plot_gradient(cancer['w1'], cancer['w2'], cancer['error'],\n",
    "                    cancer['dw1'], cancer['dw2'], scale=2)\n",
    "add_solution_path(fig, errors, path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34e928fa",
   "metadata": {},
   "source": [
    "Momentum can also be applied to the non-convex sine regression problem. Set `momentum=0` to compare with plain gradient descent at the same learning rate.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2640a59",
   "metadata": {},
   "outputs": [],
   "source": [
    "w0 = np.array([1.5, 2.])\n",
    "path = gd_momentum(w0, grad_sine_MSE, learning_rate=.1, momentum=0.5, nepochs=50)\n",
    "\n",
    "fig_loss = make_loss_curve(path, sine_MSE)\n",
    "fig_loss.show()\n",
    "\n",
    "errors = [sine_MSE(w) for w in path]\n",
    "fig = plot_gradient(sine['w0'], sine['w1'], sine['error'],\n",
    "                    sine['dw0'], sine['dw1'], scale=.1, names=names)\n",
    "add_solution_path(fig, errors, path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9f5c9dd",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "718b4a7a",
   "metadata": {},
   "source": [
    "<link rel=\"stylesheet\" href=\"berkeley.css\">\n",
    "\n",
    "<h2 class=\"cal cal-h2\">(Lecture Demo 3) Adaptive Learning Rates</h2>\n",
    "\n",
    "The optimal learning rate depends on the local curvature, which can differ between directions in weight space. Adaptive methods maintain a separate effective learning rate for each parameter $w_i$ (Bishop, Section 7.3.3). Let $g_i = \\partial E / \\partial w_i$ denote the current gradient component.\n",
    "\n",
    "**AdaGrad** accumulates the sum of squared gradients (Eqs. 7.39 and 7.40):\n",
    "$$\n",
    "r_i^{(\\tau)} = r_i^{(\\tau-1)} + g_i^2, \\qquad\n",
    "w_i^{(\\tau)} = w_i^{(\\tau-1)} - \\frac{\\eta}{\\sqrt{r_i^{(\\tau)}} + \\delta}\\, g_i.\n",
    "$$\n",
    "AdaGrad therefore takes smaller steps along parameters whose gradients have been large in the past. Because $r_i$ only grows, the effective learning rate decreases monotonically and can become too small late in training.\n",
    "\n",
    "**RMSProp** replaces the sum with an exponentially weighted moving average (Eqs. 7.41 and 7.42):\n",
    "$$\n",
    "r_i^{(\\tau)} = \\beta r_i^{(\\tau-1)} + (1 - \\beta) g_i^2, \\qquad\n",
    "w_i^{(\\tau)} = w_i^{(\\tau-1)} - \\frac{\\eta}{\\sqrt{r_i^{(\\tau)}} + \\delta}\\, g_i.\n",
    "$$\n",
    "\n",
    "**Adam** combines RMSProp with a moving average of the gradient itself, which acts as momentum (Eqs. 7.43 to 7.47):\n",
    "\\begin{align*}\n",
    "s_i^{(\\tau)} &= \\beta_1 s_i^{(\\tau-1)} + (1 - \\beta_1) g_i, &\n",
    "r_i^{(\\tau)} &= \\beta_2 r_i^{(\\tau-1)} + (1 - \\beta_2) g_i^2, \\\\\n",
    "\\widehat{s}_i^{(\\tau)} &= \\frac{s_i^{(\\tau)}}{1 - \\beta_1^\\tau}, &\n",
    "\\widehat{r}_i^{(\\tau)} &= \\frac{r_i^{(\\tau)}}{1 - \\beta_2^\\tau}, \\\\\n",
    "w_i^{(\\tau)} &= w_i^{(\\tau-1)} - \\eta \\frac{\\widehat{s}_i^{(\\tau)}}{\\sqrt{\\widehat{r}_i^{(\\tau)}} + \\delta}.\n",
    "\\end{align*}\n",
    "The factors $1/(1 - \\beta^\\tau)$ correct the bias towards zero caused by initializing $s_i^{(0)} = r_i^{(0)} = 0$. In all three methods $\\delta$ is a small constant for numerical stability. Bishop gives $\\beta_2 = 0.99$ as a typical value. The original Adam paper (Kingma and Ba, 2014) and the PyTorch default use $\\beta_1 = 0.9$ and $\\beta_2 = 0.999$, which we use below.\n",
    "\n",
    "Note that Adam's first moment $s$ is an exponential moving average: the weights $(1 - \\beta_1)\\beta_1^k$ on past gradients sum to one, so for a constant gradient $s$ approaches $\\nabla E$ itself. Adam also divides by $\\sqrt{\\widehat{r}_i}$, so its step in each coordinate has magnitude of about $\\eta$ whatever the scale of the gradient; the first step is exactly $-\\eta\\, \\mathrm{sign}(g_i)$ up to $\\delta$.\n",
    "\n",
    "*Lecture:* run from the slide \"Demo: Adam\".\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bc32ea4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def gd_adam(w_0, gradient, learning_rate=lambda t: 0.1, nepochs=100,\n",
    "            delta=1e-8, beta1=0.9, beta2=0.999):\n",
    "    \"\"\"Batch gradient descent with the Adam update.\n",
    "\n",
    "    Args:\n",
    "        w_0: Initial weight vector.\n",
    "        gradient: Function that returns the gradient of the error at w.\n",
    "        learning_rate: Function of the number of completed steps (0, 1, 2, ...) that returns eta.\n",
    "        nepochs: Number of iterations.\n",
    "        delta: Small constant added to the denominator for numerical stability.\n",
    "        beta1: Decay rate of the moving average of the gradient (first moment).\n",
    "        beta2: Decay rate of the moving average of the squared gradient (second moment).\n",
    "    Returns:\n",
    "        path: Array containing the weight vector at each iteration.\n",
    "    \"\"\"\n",
    "    w = np.asarray(w_0, dtype=float)\n",
    "    path = [w]\n",
    "    s = np.zeros_like(w)  # first moment (momentum)\n",
    "    r = np.zeros_like(w)  # second moment (RMSProp)\n",
    "    for t in range(1, nepochs + 1):\n",
    "        g = gradient(w)\n",
    "        s = beta1 * s + (1 - beta1) * g\n",
    "        r = beta2 * r + (1 - beta2) * g**2\n",
    "        s_hat = s / (1 - beta1**t)  # bias correction\n",
    "        r_hat = r / (1 - beta2**t)\n",
    "        w = w - learning_rate(t - 1) * s_hat / (np.sqrt(r_hat) + delta)  # eta^(tau-1), as in the other methods\n",
    "        path.append(w)\n",
    "    return np.array(path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "376a10f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "w0 = np.array([-1, 2.5])\n",
    "path = gd_adam(w0, grad_sine_MSE, learning_rate=lambda t: 0.1, nepochs=100)\n",
    "\n",
    "fig_loss = make_loss_curve(path, sine_MSE)\n",
    "fig_loss.show()\n",
    "\n",
    "errors = [sine_MSE(w) for w in path]\n",
    "fig = plot_gradient(sine['w0'], sine['w1'], sine['error'],\n",
    "                    sine['dw0'], sine['dw1'], scale=.1, names=names)\n",
    "add_solution_path(fig, errors, path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1160691",
   "metadata": {},
   "source": [
    "<link rel=\"stylesheet\" href=\"berkeley.css\">\n",
    "\n",
    "<h3 class=\"cal cal-h3\">Separating the Update Rule from the Optimization Loop</h3>\n",
    "\n",
    "The algorithms above differ only in how the gradient is turned into an update. We therefore write each update rule as an optimizer object with a `step(w, g)` method, which returns the new weight vector given the current weights and gradient, and write a single optimization loop. Each optimizer takes the learning rate as a function of the number of completed steps $t = 0, 1, 2, \\dots$, so the step that produces $w^{(\\tau)}$ uses $\\eta^{(\\tau-1)}$ (Bishop, Eq. 7.35), as `torch.optim.lr_scheduler` does. `MomentumOptimizer` applies the current learning rate to the new gradient only, as in Bishop's form; PyTorch instead multiplies the whole velocity by the current learning rate, so the two differ while the learning rate is changing. The same optimizer objects are used with stochastic gradient descent below. This mirrors the design of `torch.optim`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f044f5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "class VanillaOptimizer:\n",
    "    \"\"\"Gradient descent update w <- w - eta_t g.\"\"\"\n",
    "    def __init__(self, learning_rate=lambda t: 0.1):\n",
    "        self.learning_rate = learning_rate\n",
    "        self.t = 0\n",
    "\n",
    "    def step(self, w, g):\n",
    "        w = w - self.learning_rate(self.t) * g\n",
    "        self.t += 1\n",
    "        return w\n",
    "\n",
    "\n",
    "class MomentumOptimizer:\n",
    "    \"\"\"Momentum: v <- mu v - eta_t g, w <- w + v (v stores the previous update Delta w).\"\"\"\n",
    "    def __init__(self, learning_rate=lambda t: 0.1, momentum=0.9):\n",
    "        self.learning_rate = learning_rate\n",
    "        self.momentum = momentum\n",
    "        self.t = 0\n",
    "        self.v = None\n",
    "\n",
    "    def step(self, w, g):\n",
    "        if self.v is None:\n",
    "            self.v = np.zeros_like(w)\n",
    "        self.v = self.momentum * self.v - self.learning_rate(self.t) * g\n",
    "        self.t += 1\n",
    "        return w + self.v\n",
    "\n",
    "\n",
    "class AdaGradOptimizer:\n",
    "    \"\"\"AdaGrad: r <- r + g^2, w <- w - eta_t g / (sqrt(r) + delta).\"\"\"\n",
    "    def __init__(self, learning_rate=lambda t: 0.1, delta=1e-8):\n",
    "        self.learning_rate = learning_rate\n",
    "        self.delta = delta\n",
    "        self.t = 0\n",
    "        self.r = None\n",
    "\n",
    "    def step(self, w, g):\n",
    "        if self.r is None:\n",
    "            self.r = np.zeros_like(w)\n",
    "        self.r = self.r + g**2\n",
    "        w = w - self.learning_rate(self.t) * g / (np.sqrt(self.r) + self.delta)\n",
    "        self.t += 1\n",
    "        return w\n",
    "\n",
    "\n",
    "class RMSPropOptimizer:\n",
    "    \"\"\"RMSProp: r <- beta r + (1 - beta) g^2, w <- w - eta_t g / (sqrt(r) + delta).\"\"\"\n",
    "    def __init__(self, learning_rate=lambda t: 0.01, beta=0.9, delta=1e-8):\n",
    "        self.learning_rate = learning_rate\n",
    "        self.beta = beta\n",
    "        self.delta = delta\n",
    "        self.t = 0\n",
    "        self.r = None\n",
    "\n",
    "    def step(self, w, g):\n",
    "        if self.r is None:\n",
    "            self.r = np.zeros_like(w)\n",
    "        self.r = self.beta * self.r + (1 - self.beta) * g**2\n",
    "        w = w - self.learning_rate(self.t) * g / (np.sqrt(self.r) + self.delta)\n",
    "        self.t += 1\n",
    "        return w\n",
    "\n",
    "\n",
    "class AdamOptimizer:\n",
    "    \"\"\"Adam with bias-corrected first and second moment estimates.\"\"\"\n",
    "    def __init__(self, learning_rate=lambda t: 0.1, beta1=0.9, beta2=0.999, delta=1e-8):\n",
    "        self.learning_rate = learning_rate\n",
    "        self.beta1 = beta1\n",
    "        self.beta2 = beta2\n",
    "        self.delta = delta\n",
    "        self.t = 0\n",
    "        self.s = None\n",
    "        self.r = None\n",
    "\n",
    "    def step(self, w, g):\n",
    "        if self.s is None:\n",
    "            self.s = np.zeros_like(w)  # first moment (momentum)\n",
    "            self.r = np.zeros_like(w)  # second moment (RMSProp)\n",
    "        eta = self.learning_rate(self.t)  # schedule at the number of completed steps\n",
    "        self.t += 1\n",
    "        self.s = self.beta1 * self.s + (1 - self.beta1) * g\n",
    "        self.r = self.beta2 * self.r + (1 - self.beta2) * g**2\n",
    "        s_hat = self.s / (1 - self.beta1**self.t)  # bias correction\n",
    "        r_hat = self.r / (1 - self.beta2**self.t)\n",
    "        return w - eta * s_hat / (np.sqrt(r_hat) + self.delta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14788759",
   "metadata": {},
   "outputs": [],
   "source": [
    "def gd(w_0, gradient, optimizer, nepochs=100, epsilon=1e-8):\n",
    "    \"\"\"Batch gradient descent with a pluggable update rule.\n",
    "\n",
    "    Args:\n",
    "        w_0: Initial weight vector.\n",
    "        gradient: Function that returns the gradient of the error at w.\n",
    "        optimizer: Object with a step(w, g) method that returns the updated weights.\n",
    "        nepochs: Maximum number of iterations.\n",
    "        epsilon: Stop when the norm of the update falls below this threshold.\n",
    "    Returns:\n",
    "        path: Array containing the weight vector at each iteration.\n",
    "    \"\"\"\n",
    "    w = np.asarray(w_0, dtype=float)\n",
    "    path = [w]\n",
    "    for _ in range(nepochs):\n",
    "        w = optimizer.step(w, gradient(w))\n",
    "        path.append(w)\n",
    "        if np.linalg.norm(w - path[-2]) < epsilon:\n",
    "            break\n",
    "    return np.array(path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84dc9898",
   "metadata": {},
   "source": [
    "The loss curves below compare the five update rules on the sine regression problem from the same initial point. The learning rates are chosen separately for each method, because the scale of the update differs between methods.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "852d08a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "w0 = np.array([-1, 2.5])\n",
    "optimizers = {\n",
    "    'Gradient descent': VanillaOptimizer(learning_rate=lambda t: 0.1),\n",
    "    'Momentum': MomentumOptimizer(learning_rate=lambda t: 0.01, momentum=0.9),\n",
    "    'AdaGrad': AdaGradOptimizer(learning_rate=lambda t: 0.1),\n",
    "    'RMSProp': RMSPropOptimizer(learning_rate=lambda t: 0.1),\n",
    "    'Adam': AdamOptimizer(learning_rate=lambda t: 0.1),\n",
    "}\n",
    "fig = go.Figure()\n",
    "paths = {}\n",
    "for name, opt in optimizers.items():\n",
    "    paths[name] = gd(w0, grad_sine_MSE, opt, nepochs=100)\n",
    "    fig.add_scatter(y=[sine_MSE(w) for w in paths[name]], mode='lines', name=name,\n",
    "                    line=dict(width=3))\n",
    "fig.update_layout(xaxis_title='Iteration', yaxis_title='Error',\n",
    "                  margin=dict(l=20, r=20, t=20, b=20))\n",
    "fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c680b03",
   "metadata": {},
   "outputs": [],
   "source": [
    "path = paths['Adam']\n",
    "errors = [sine_MSE(w) for w in path]\n",
    "fig = plot_gradient(sine['w0'], sine['w1'], sine['error'],\n",
    "                    sine['dw0'], sine['dw1'], scale=.1, names=names)\n",
    "add_solution_path(fig, errors, path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85273b49",
   "metadata": {},
   "source": [
    "<link rel=\"stylesheet\" href=\"berkeley.css\">\n",
    "\n",
    "<h2 class=\"cal cal-h2\">(Lecture Demo 4) Learning Rate Schedules</h2>\n",
    "\n",
    "In practice a larger learning rate is used at the start of training and is then reduced, so that $\\eta$ becomes a function of the iteration index $\\tau$ (Bishop, Eq. 7.35):\n",
    "$$\n",
    "w^{(\\tau)} = w^{(\\tau-1)} - \\eta^{(\\tau-1)} \\nabla E\\big(w^{(\\tau-1)}\\big).\n",
    "$$\n",
    "Each schedule below is implemented as a function that returns a callable $\\tau \\mapsto \\eta^{(\\tau)}$. Here $\\tau$ counts iterations (gradient steps), and $T$ is the total number of iterations.\n",
    "\n",
    "*Lecture:* the comparison at the end of this section is the figure on \"Demo: Learning Rate Schedules\".\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ac5984d",
   "metadata": {},
   "source": [
    "<link rel=\"stylesheet\" href=\"berkeley.css\">\n",
    "\n",
    "<h3 class=\"cal cal-h3\">Linear Schedule</h3>\n",
    "\n",
    "$$\n",
    "\\eta^{(\\tau)} = \\left(1 - \\frac{\\tau}{T}\\right) \\eta_{\\text{start}} + \\frac{\\tau}{T} \\eta_{\\text{end}},\n",
    "$$\n",
    "after which the learning rate is held at $\\eta_{\\text{end}}$ (Bishop, Eq. 7.36).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9ba3ba5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def linear_learning_rate(initial_lr, T, end_lr=0):\n",
    "    \"\"\"Linear decay from initial_lr to end_lr over T iterations, then constant.\"\"\"\n",
    "    return lambda t: initial_lr * (1 - min(1, t / T)) + end_lr * min(1, t / T)\n",
    "\n",
    "def plot_schedule(lr, T=100):\n",
    "    fig = px.line(x=np.arange(T), y=[lr(t) for t in range(T)], markers=True,\n",
    "                  labels={'x': 'Iteration', 'y': 'Learning Rate'})\n",
    "    fig.update_traces(line_width=4)\n",
    "    fig.update_layout(margin=dict(l=20, r=20, t=20, b=20))\n",
    "    return fig\n",
    "\n",
    "plot_schedule(linear_learning_rate(initial_lr=0.1, T=100, end_lr=0.01))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee1073ce",
   "metadata": {},
   "source": [
    "<link rel=\"stylesheet\" href=\"berkeley.css\">\n",
    "\n",
    "<h3 class=\"cal cal-h3\">Power-Law Schedule</h3>\n",
    "\n",
    "$$\n",
    "\\eta^{(\\tau)} = \\eta_{\\text{start}} \\left(1 + \\frac{\\tau}{s}\\right)^{-c}\n",
    "$$\n",
    "(Bishop, Eq. 7.37), where $s$ sets the time scale of the decay and $c$ its exponent.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80a86a59",
   "metadata": {},
   "outputs": [],
   "source": [
    "def power_law_learning_rate(initial_lr, s=1, c=1):\n",
    "    \"\"\"Power-law decay: initial_lr * (1 + t/s)^(-c).\"\"\"\n",
    "    return lambda t: initial_lr * (1 + t / s) ** (-c)\n",
    "\n",
    "plot_schedule(power_law_learning_rate(initial_lr=0.1, s=10, c=1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08a4b25a",
   "metadata": {},
   "source": [
    "<link rel=\"stylesheet\" href=\"berkeley.css\">\n",
    "\n",
    "<h3 class=\"cal cal-h3\">Exponential and Step Decay</h3>\n",
    "\n",
    "The exponential schedule (Bishop, Eq. 7.38) is\n",
    "$$\n",
    "\\eta^{(\\tau)} = \\eta_{\\text{start}}\\, c^{\\tau/s}, \\qquad 0 < c < 1.\n",
    "$$\n",
    "The slides use the common variant **step decay**, which holds the learning rate constant for $s$ iterations at a time: $\\eta^{(\\tau)} = \\eta_{\\text{start}}\\, c^{\\lfloor \\tau/s \\rfloor}$.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e65098b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def exponential_learning_rate(initial_lr, c=0.9, s=10, staircase=False):\n",
    "    \"\"\"Exponential decay initial_lr * c^(t/s), or step decay c^floor(t/s) if staircase.\"\"\"\n",
    "    if staircase:\n",
    "        return lambda t: initial_lr * c ** np.floor(t / s)\n",
    "    return lambda t: initial_lr * c ** (t / s)\n",
    "\n",
    "plot_schedule(exponential_learning_rate(initial_lr=0.1, c=0.9, s=10, staircase=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95e6c13e",
   "metadata": {},
   "source": [
    "<link rel=\"stylesheet\" href=\"berkeley.css\">\n",
    "\n",
    "<h3 class=\"cal cal-h3\">Cosine Schedule with Linear Warmup</h3>\n",
    "\n",
    "The cosine schedule decays from $\\eta_{\\max}$ to $\\eta_{\\min}$ over $T$ iterations:\n",
    "$$\n",
    "\\eta^{(\\tau)} = \\eta_{\\min} + \\frac{1}{2}\\left(\\eta_{\\max} - \\eta_{\\min}\\right)\\left(1 + \\cos\\left(\\frac{\\pi \\tau}{T}\\right)\\right).\n",
    "$$\n",
    "Large models are commonly trained with a short **warmup** phase before the decay, during which the learning rate increases linearly from 0 to $\\eta_{\\max}$ over the first $T_w$ iterations, $\\eta^{(\\tau)} = \\eta_{\\max}\\, \\tau / T_w$ for $\\tau < T_w$, and then follows the cosine schedule over the remaining $T - T_w$ iterations. Warmup avoids large, poorly directed steps at the start of training, when the optimizer state (for example the moving averages in Adam) has not yet adapted.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54e7f225",
   "metadata": {},
   "outputs": [],
   "source": [
    "def cosine_learning_rate(T, max_lr, min_lr=0):\n",
    "    \"\"\"Cosine decay from max_lr to min_lr over T iterations.\"\"\"\n",
    "    return lambda t: min_lr + 0.5 * (max_lr - min_lr) * (1 + np.cos(np.pi * min(t, T) / T))\n",
    "\n",
    "def warmup_cosine_learning_rate(T, T_warmup, max_lr, min_lr=0):\n",
    "    \"\"\"Linear warmup over T_warmup iterations followed by cosine decay to min_lr at T.\"\"\"\n",
    "    cosine = cosine_learning_rate(T - T_warmup, max_lr, min_lr)\n",
    "    return lambda t: max_lr * t / T_warmup if t < T_warmup else cosine(t - T_warmup)\n",
    "\n",
    "plot_schedule(warmup_cosine_learning_rate(T=100, T_warmup=10, max_lr=0.1, min_lr=0.01))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67895d25",
   "metadata": {},
   "source": [
    "<link rel=\"stylesheet\" href=\"berkeley.css\">\n",
    "\n",
    "<h3 class=\"cal cal-h3\">Comparing Schedules</h3>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a14803fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "T = 100\n",
    "schedules = {\n",
    "    'Linear': linear_learning_rate(initial_lr=0.1, T=T, end_lr=0.01),\n",
    "    'Power law': power_law_learning_rate(initial_lr=0.1, s=10, c=1),\n",
    "    'Step decay': exponential_learning_rate(initial_lr=0.1, c=0.9, s=5, staircase=True),\n",
    "    'Cosine': cosine_learning_rate(T=T, max_lr=0.1, min_lr=0.01),\n",
    "    'Warmup + cosine': warmup_cosine_learning_rate(T=T, T_warmup=10, max_lr=0.1, min_lr=0.01),\n",
    "}\n",
    "t = np.arange(T)\n",
    "fig = go.Figure()\n",
    "for name, lr in schedules.items():\n",
    "    fig.add_scatter(x=t, y=[lr(tau) for tau in t], mode='lines+markers', name=name,\n",
    "                    line=dict(width=4))\n",
    "fig.update_layout(margin=dict(l=20, r=20, t=20, b=20),\n",
    "                  xaxis_title='Iteration', yaxis_title='Learning Rate', height=600)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "160a6b1d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "def76314",
   "metadata": {},
   "source": [
    "<link rel=\"stylesheet\" href=\"berkeley.css\">\n",
    "\n",
    "<h2 class=\"cal cal-h2\">(Lecture Demo 5) Stochastic Gradient Descent</h2>\n",
    "\n",
    "The error function is an average over the training data,\n",
    "$$\n",
    "E(w) = \\frac{1}{N} \\sum_{n=1}^N E_n(w),\n",
    "$$\n",
    "so each batch gradient step evaluates $N$ per-example gradients. For a model with $W$ parameters each of these costs $O(W)$, so a batch step costs $O(NW)$. Stochastic gradient descent (Bishop, Section 7.2.3) replaces $\\nabla E$ with the gradient of the error on a single data point $n$:\n",
    "$$\n",
    "w^{(\\tau)} = w^{(\\tau-1)} - \\eta \\nabla E_n\\big(w^{(\\tau-1)}\\big).\n",
    "$$\n",
    "If $n$ is drawn uniformly from $\\{1, \\dots, N\\}$, the stochastic gradient is an **unbiased** estimate of the batch gradient:\n",
    "$$\n",
    "\\mathbb{E}_n\\big[\\nabla E_n(w)\\big] = \\frac{1}{N} \\sum_{n=1}^N \\nabla E_n(w) = \\nabla E(w).\n",
    "$$\n",
    "Bishop defines $E(w) = \\sum_n E_n(w)$ as a sum rather than an average. The two conventions differ by the constant factor $N$, which is absorbed into the learning rate.\n",
    "\n",
    "The slide \"Stochastic Gradient Descent\" samples a random index at each step. In practice the data are shuffled at the start of each epoch and then visited in order, so each point is used exactly once per epoch (\"Stochastic Gradient Descent (Shuffling)\"). The implementation below runs for a fixed number of epochs and uses a seeded random number generator so that the results are reproducible.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8260250",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sgd(w_0, x, t, gradient, optimizer, nepochs=2, seed=0):\n",
    "    \"\"\"Stochastic gradient descent with one data point per step.\n",
    "\n",
    "    Args:\n",
    "        w_0: Initial weight vector.\n",
    "        x, t: Training inputs and targets.\n",
    "        gradient: Function gradient(w, x_n, t_n) returning the gradient on the given data.\n",
    "        optimizer: Object with a step(w, g) method that returns the updated weights.\n",
    "        nepochs: Number of passes over the data.\n",
    "        seed: Seed for the random shuffling of the data.\n",
    "    Returns:\n",
    "        path: Array containing the weight vector after each step.\n",
    "    \"\"\"\n",
    "    rng = np.random.default_rng(seed)\n",
    "    w = np.asarray(w_0, dtype=float)\n",
    "    path = [w]\n",
    "    for _ in range(nepochs):\n",
    "        for n in rng.permutation(len(t)):  # reshuffle every epoch\n",
    "            w = optimizer.step(w, gradient(w, x[n], t[n]))\n",
    "            path.append(w)\n",
    "    return np.array(path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9be3373a",
   "metadata": {},
   "source": [
    "SGD with Adam on the sine regression problem. Each epoch consists of $N = 200$ steps, each using a single data point. The path is noisy because each step follows the gradient of a single term.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8ffc77a",
   "metadata": {},
   "outputs": [],
   "source": [
    "w0 = np.array([-1, 2.5])\n",
    "path = sgd(w0, sine['x'], sine['y'], grad_sine_MSE,\n",
    "           AdamOptimizer(learning_rate=lambda t: 0.1), nepochs=2)\n",
    "\n",
    "errors = [sine_MSE(w) for w in path]\n",
    "fig = plot_gradient(sine['w0'], sine['w1'], sine['error'],\n",
    "                    sine['dw0'], sine['dw1'], scale=.1, names=names)\n",
    "fig = add_solution_path(fig, errors, path)\n",
    "fig.show()\n",
    "\n",
    "fig_loss = make_loss_curve(path, sine_MSE)\n",
    "fig_loss.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a8d1df5",
   "metadata": {},
   "source": [
    "SGD with Adam and a cosine learning rate schedule on the logistic regression problem, for one epoch of $N = 569$ steps. The decreasing learning rate reduces the noise in the iterates as training proceeds. This is the figure on \"Demo: Stochastic Gradient Descent\".\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "630dfe1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "w0 = np.array([-1., 2.])\n",
    "nepochs = 1\n",
    "coslr = cosine_learning_rate(T=nepochs * len(cancer['t']), max_lr=0.1, min_lr=0.01)\n",
    "path = sgd(w0, cancer['x'], cancer['t'], grad_NLL,\n",
    "           AdamOptimizer(learning_rate=coslr), nepochs=nepochs)\n",
    "\n",
    "errors = [neg_log_likelihood(w) for w in path]\n",
    "fig = plot_gradient(cancer['w1'], cancer['w2'], cancer['error'],\n",
    "                    cancer['dw1'], cancer['dw2'], scale=1)\n",
    "fig = add_solution_path(fig, errors, path)\n",
    "fig.show()\n",
    "\n",
    "fig_loss = make_loss_curve(path, neg_log_likelihood)\n",
    "fig_loss.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcd38c03",
   "metadata": {},
   "source": [
    "<link rel=\"stylesheet\" href=\"berkeley.css\">\n",
    "\n",
    "<h2 class=\"cal cal-h2\">(Lecture Demo 6) Mini-Batch Stochastic Gradient Descent</h2>\n",
    "\n",
    "A mini-batch $\\mathcal{B}$ of $B$ data points gives the gradient estimate\n",
    "$$\n",
    "\\nabla E_{\\mathcal{B}}(w) = \\frac{1}{B} \\sum_{n \\in \\mathcal{B}} \\nabla E_n(w),\n",
    "$$\n",
    "which is also unbiased and costs $O(BW)$. For $B = 1$ this is SGD, and for $B = N$ it is batch gradient descent. If $\\sigma$ is the standard deviation of a single-point gradient component, the standard deviation of the mini-batch estimate is approximately $\\sigma/\\sqrt{B}$ (Bishop, Section 7.2.4 and Exercise 7.8). Increasing the batch size by a factor of 100 therefore reduces the noise only by a factor of 10.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33c945d4",
   "metadata": {},
   "source": [
    "Mini-batch SGD shuffles the data at the start of each epoch and then takes successive blocks of $B$ points (Bishop, Algorithm 7.2). Each epoch therefore has $\\lceil N/B \\rceil$ steps (for the Slido question, $N/2$ steps when $B = 2$ and $N$ is even). With $N = 200$ and $B = 32$ each epoch has $\\lceil 200/32 \\rceil = 7$ steps, the last of which uses the remaining 8 points.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "544c2111",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mb_sgd(w_0, x, t, gradient, optimizer, batch_size=32, nepochs=2, seed=0):\n",
    "    \"\"\"Mini-batch stochastic gradient descent.\n",
    "\n",
    "    Args:\n",
    "        w_0: Initial weight vector.\n",
    "        x, t: Training inputs and targets.\n",
    "        gradient: Function gradient(w, x_batch, t_batch) returning the average gradient on the batch.\n",
    "        optimizer: Object with a step(w, g) method that returns the updated weights.\n",
    "        batch_size: Number of data points B in each mini-batch.\n",
    "        nepochs: Number of passes over the data.\n",
    "        seed: Seed for the random shuffling of the data.\n",
    "    Returns:\n",
    "        path: Array containing the weight vector after each step.\n",
    "    \"\"\"\n",
    "    rng = np.random.default_rng(seed)\n",
    "    w = np.asarray(w_0, dtype=float)\n",
    "    path = [w]\n",
    "    for _ in range(nepochs):\n",
    "        order = rng.permutation(len(t))  # reshuffle every epoch\n",
    "        for start in range(0, len(t), batch_size):\n",
    "            batch = order[start:start + batch_size]\n",
    "            w = optimizer.step(w, gradient(w, x[batch], t[batch]))\n",
    "            path.append(w)\n",
    "    return np.array(path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3264945",
   "metadata": {},
   "outputs": [],
   "source": [
    "w0 = np.array([-1, 2.5])\n",
    "path = mb_sgd(w0, sine['x'], sine['y'], grad_sine_MSE,\n",
    "              AdamOptimizer(learning_rate=lambda t: 0.1),\n",
    "              batch_size=32, nepochs=10)\n",
    "\n",
    "errors = [sine_MSE(w) for w in path]\n",
    "fig = plot_gradient(sine['w0'], sine['w1'], sine['error'],\n",
    "                    sine['dw0'], sine['dw1'], scale=.1, names=names)\n",
    "fig = add_solution_path(fig, errors, path)\n",
    "fig.show()\n",
    "\n",
    "fig_loss = make_loss_curve(path, sine_MSE)\n",
    "fig_loss.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9a2fc32",
   "metadata": {},
   "source": [
    "Mini-batch SGD with Adam on the logistic regression problem, with $B = 32$ and 5 epochs. This is the figure on \"Demo: Mini-batch Gradient Descent\".\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2531b344",
   "metadata": {},
   "outputs": [],
   "source": [
    "w0 = np.array([0., -4.])\n",
    "path = mb_sgd(w0, cancer['x'], cancer['t'], grad_NLL,\n",
    "              AdamOptimizer(learning_rate=lambda t: 0.2),\n",
    "              batch_size=32, nepochs=5)\n",
    "\n",
    "errors = [neg_log_likelihood(w) for w in path]\n",
    "fig = plot_gradient(cancer['w1'], cancer['w2'], cancer['error'],\n",
    "                    cancer['dw1'], cancer['dw2'], scale=2)\n",
    "fig = add_solution_path(fig, errors, path)\n",
    "fig.show()\n",
    "\n",
    "fig_loss = make_loss_curve(path, neg_log_likelihood)\n",
    "fig_loss.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py_3_11",
   "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.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
