{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cell-000",
   "metadata": {},
   "source": [
    "---\n",
    "title: Checking Manual Gradients Against Autodiff on MNIST\n",
    "category: Deep Learning\n",
    "tags:\n",
    "  - deep learning\n",
    "  - autodiff\n",
    "  - gradients\n",
    "  - mnist\n",
    "  - pytorch\n",
    "date: 2026-07-31\n",
    "description: Compare a hand-derived softmax regression gradient against PyTorch autograd on a real MNIST minibatch.\n",
    "notebook_page: false\n",
    "aliases:\n",
    "  - /deep-learning/checking-manual-gradients-against-autodiff-on-mnist\n",
    "sources:\n",
    "  - https://huggingface.co/datasets/ylecun/mnist\n",
    "  - https://docs.pytorch.org/tutorials/beginner/basics/autogradqs_tutorial.html\n",
    "---\n",
    "\n",
    "## Goal\n",
    "\n",
    "The shortest path to trusting autodiff is to catch it agreeing with a derivation you can inspect. We will use a real minibatch from [MNIST](https://huggingface.co/datasets/ylecun/mnist), derive the gradient of softmax regression by hand, and compare it against `torch.autograd`.\n",
    "\n",
    "For logits \\(Z = XW + b\\), probabilities \\(P = \\\\text{softmax}(Z)\\), and one-hot targets \\(Y\\), the cross-entropy derivative is\n",
    "\n",
    "$$\n",
    "\\\\frac{\\\\partial \\\\mathcal{L}}{\\\\partial Z} = \\\\frac{P - Y}{B}.\n",
    "$$\n",
    "\n",
    "The weight gradient follows immediately:\n",
    "\n",
    "$$\n",
    "\\\\frac{\\\\partial \\\\mathcal{L}}{\\\\partial W} = X^\\\\top \\\\frac{P - Y}{B}.\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "cell-001",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-31T21:58:42.005159Z",
     "iopub.status.busy": "2026-07-31T21:58:42.005070Z",
     "iopub.status.idle": "2026-07-31T21:58:43.291783Z",
     "shell.execute_reply": "2026-07-31T21:58:43.291137Z"
    }
   },
   "outputs": [],
   "source": [
    "## %pip install -q datasets torch numpy matplotlib\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from datasets import load_dataset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cell-002",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-31T21:58:43.293106Z",
     "iopub.status.busy": "2026-07-31T21:58:43.292957Z",
     "iopub.status.idle": "2026-07-31T21:58:44.918437Z",
     "shell.execute_reply": "2026-07-31T21:58:44.917908Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(256, 784) (256,)\n"
     ]
    }
   ],
   "source": [
    "ds = load_dataset(\"ylecun/mnist\")\n",
    "train_subset = ds[\"train\"].select(range(1024))\n",
    "test_subset = ds[\"test\"].select(range(1024))\n",
    "\n",
    "x = np.stack([np.array(example[\"image\"], dtype=np.float32).reshape(-1) for example in train_subset]) / 255.0\n",
    "y = np.array(train_subset[\"label\"], dtype=np.int64)\n",
    "\n",
    "batch_x = x[:256]\n",
    "batch_y = y[:256]\n",
    "\n",
    "print(batch_x.shape, batch_y.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cell-003",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-31T21:58:44.919745Z",
     "iopub.status.busy": "2026-07-31T21:58:44.919625Z",
     "iopub.status.idle": "2026-07-31T21:58:44.922717Z",
     "shell.execute_reply": "2026-07-31T21:58:44.922167Z"
    }
   },
   "outputs": [],
   "source": [
    "def softmax_numpy(logits: np.ndarray) -> np.ndarray:\n",
    "    shifted = logits - logits.max(axis=1, keepdims=True)\n",
    "    exp = np.exp(shifted)\n",
    "    return exp / exp.sum(axis=1, keepdims=True)\n",
    "\n",
    "\n",
    "def manual_gradients(xb: np.ndarray, yb: np.ndarray, W: np.ndarray, b: np.ndarray):\n",
    "    logits = xb @ W + b\n",
    "    probs = softmax_numpy(logits)\n",
    "    one_hot = np.zeros_like(probs)\n",
    "    one_hot[np.arange(len(yb)), yb] = 1.0\n",
    "    dz = (probs - one_hot) / len(yb)\n",
    "    dW = xb.T @ dz\n",
    "    db = dz.sum(axis=0, keepdims=True)\n",
    "    loss = -np.log(probs[np.arange(len(yb)), yb] + 1e-9).mean()\n",
    "    return loss, dW, db\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cell-004",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-31T21:58:44.923894Z",
     "iopub.status.busy": "2026-07-31T21:58:44.923798Z",
     "iopub.status.idle": "2026-07-31T21:58:44.945783Z",
     "shell.execute_reply": "2026-07-31T21:58:44.945320Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "manual loss: 2.3231063\n",
      "autodiff loss: 2.323106527328491\n",
      "max |dW_manual - dW_autodiff| = 1.4901161e-08\n",
      "max |db_manual - db_autodiff| = 1.8626451e-08\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(7)\n",
    "\n",
    "W_np = np.random.randn(784, 10).astype(np.float32) * 0.01\n",
    "b_np = np.zeros((1, 10), dtype=np.float32)\n",
    "\n",
    "manual_loss, manual_dW, manual_db = manual_gradients(batch_x, batch_y, W_np, b_np)\n",
    "\n",
    "xb_t = torch.tensor(batch_x, dtype=torch.float32)\n",
    "yb_t = torch.tensor(batch_y, dtype=torch.long)\n",
    "W_t = torch.tensor(W_np, dtype=torch.float32, requires_grad=True)\n",
    "b_t = torch.tensor(b_np, dtype=torch.float32, requires_grad=True)\n",
    "\n",
    "logits_t = xb_t @ W_t + b_t\n",
    "loss_t = F.cross_entropy(logits_t, yb_t)\n",
    "loss_t.backward()\n",
    "\n",
    "max_abs_dW = np.abs(manual_dW - W_t.grad.detach().numpy()).max()\n",
    "max_abs_db = np.abs(manual_db - b_t.grad.detach().numpy()).max()\n",
    "\n",
    "print(\"manual loss:\", manual_loss)\n",
    "print(\"autodiff loss:\", float(loss_t))\n",
    "print(\"max |dW_manual - dW_autodiff| =\", max_abs_dW)\n",
    "print(\"max |db_manual - db_autodiff| =\", max_abs_db)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-005",
   "metadata": {},
   "source": [
    "## Finite Difference Sanity Check\n",
    "\n",
    "Autodiff and the symbolic derivation can still agree on the same bug. A cheap extra check is finite differences on a few random coordinates:\n",
    "\n",
    "$$\n",
    "\\\\frac{\\\\partial \\\\mathcal{L}}{\\\\partial \\\\theta_j} \\\\approx \\\\frac{\\\\mathcal{L}(\\\\theta_j + \\\\varepsilon) - \\\\mathcal{L}(\\\\theta_j - \\\\varepsilon)}{2\\\\varepsilon}.\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "cell-006",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-31T21:58:44.946861Z",
     "iopub.status.busy": "2026-07-31T21:58:44.946761Z",
     "iopub.status.idle": "2026-07-31T21:58:44.950182Z",
     "shell.execute_reply": "2026-07-31T21:58:44.949793Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0, 0) finite-diff = 0.0 manual = 0.0\n",
      "(123, 3) finite-diff = -0.00476837158203125 manual = -0.004906716\n",
      "(511, 7) finite-diff = 0.015497207641601562 manual = 0.01614688\n"
     ]
    }
   ],
   "source": [
    "def loss_only(xb: np.ndarray, yb: np.ndarray, W: np.ndarray, b: np.ndarray) -> float:\n",
    "    logits = xb @ W + b\n",
    "    probs = softmax_numpy(logits)\n",
    "    return float(-np.log(probs[np.arange(len(yb)), yb] + 1e-9).mean())\n",
    "\n",
    "\n",
    "eps = 1e-4\n",
    "checks = [(0, 0), (123, 3), (511, 7)]\n",
    "\n",
    "for i, j in checks:\n",
    "    W_pos = W_np.copy()\n",
    "    W_neg = W_np.copy()\n",
    "    W_pos[i, j] += eps\n",
    "    W_neg[i, j] -= eps\n",
    "    fd = (loss_only(batch_x, batch_y, W_pos, b_np) - loss_only(batch_x, batch_y, W_neg, b_np)) / (2 * eps)\n",
    "    print((i, j), \"finite-diff =\", fd, \"manual =\", manual_dW[i, j])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "cell-007",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-31T21:58:44.951250Z",
     "iopub.status.busy": "2026-07-31T21:58:44.951156Z",
     "iopub.status.idle": "2026-07-31T21:58:45.659085Z",
     "shell.execute_reply": "2026-07-31T21:58:45.658599Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch=00 train_acc=0.3193 test_acc=0.2715\n",
      "epoch=01 train_acc=0.5527 test_acc=0.4365\n",
      "epoch=02 train_acc=0.6621 test_acc=0.5186\n",
      "epoch=03 train_acc=0.7275 test_acc=0.6084\n",
      "epoch=04 train_acc=0.7646 test_acc=0.6396\n",
      "epoch=05 train_acc=0.7842 test_acc=0.6758\n",
      "epoch=06 train_acc=0.8018 test_acc=0.6895\n",
      "epoch=07 train_acc=0.8115 test_acc=0.7002\n",
      "epoch=08 train_acc=0.8125 test_acc=0.7090\n",
      "epoch=09 train_acc=0.8164 test_acc=0.7168\n"
     ]
    }
   ],
   "source": [
    "class SoftmaxRegressor(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.linear = torch.nn.Linear(784, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.linear(x)\n",
    "\n",
    "\n",
    "model = SoftmaxRegressor()\n",
    "opt = torch.optim.Adam(model.parameters(), lr=1e-3)\n",
    "\n",
    "x_train = torch.tensor(x[:1024], dtype=torch.float32)\n",
    "y_train = torch.tensor(y[:1024], dtype=torch.long)\n",
    "x_test = torch.tensor(\n",
    "    np.stack([np.array(example[\"image\"], dtype=np.float32).reshape(-1) for example in test_subset]) / 255.0,\n",
    "    dtype=torch.float32,\n",
    ")\n",
    "y_test = torch.tensor(test_subset[\"label\"], dtype=torch.long)\n",
    "\n",
    "history = []\n",
    "batch_size = 256\n",
    "\n",
    "for epoch in range(10):\n",
    "    perm = torch.randperm(len(x_train))\n",
    "    for start in range(0, len(x_train), batch_size):\n",
    "        idx = perm[start : start + batch_size]\n",
    "        logits = model(x_train[idx])\n",
    "        loss = F.cross_entropy(logits, y_train[idx])\n",
    "        opt.zero_grad()\n",
    "        loss.backward()\n",
    "        opt.step()\n",
    "\n",
    "    with torch.no_grad():\n",
    "        train_acc = (model(x_train).argmax(dim=1) == y_train).float().mean().item()\n",
    "        test_acc = (model(x_test).argmax(dim=1) == y_test).float().mean().item()\n",
    "        history.append((train_acc, test_acc))\n",
    "        print(f\"epoch={epoch:02d} train_acc={train_acc:.4f} test_acc={test_acc:.4f}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-008",
   "metadata": {},
   "source": [
    "## Takeaway\n",
    "\n",
    "The interesting result is not that `autograd` works. The interesting result is that a short derivation, finite differences, and the framework implementation all agree on the same real minibatch. Once that is in place, the rest of deep learning code becomes much less mysterious.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.12.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}