Neural Networks from Scratch

Purpose

This note is the foundation under the rest of the deep-learning section. The goal is to make the standard neural-network training loop feel mechanical rather than mysterious:

• define a computation
• define a scalar loss
• differentiate it with the chain rule
• update parameters

Rumelhart, Hinton, and Williams made the key move in 1986. Once the network is a composition of differentiable operations, the gradient can be pushed backward one local Jacobian at a time.

One Affine Layer

For input $x\in {\mathbb{R}}^{d_{in}}$, weights $W\in {\mathbb{R}}^{d_{out}\times d_{in}}$, and bias $b\in {\mathbb{R}}^{d_{out}}$:

$$z=Wx+b$$

This is just an affine map. If a network were only a stack of affine maps, the whole composition would still be affine:

$$W_2(W_{1}x+b_1)+b_2=(W_2{W}_1)x+(W_2{b}_1+b_2)$$

That is why hidden-layer nonlinearities are necessary. They prevent depth from collapsing into one linear classifier.

A Two-Layer Classifier

For hidden width $h$ and $K$ classes:

$$\begin{array}{rl}{z}_1& =W_{1}x+b_1\\ a_1& =\varphi (z_1)\\ z_2& =W_2{a}_1+b_2\\ \hat{y}& =\text{softmax}(z_2)\end{array}$$

The softmax is

$${\hat{y}}_k=\frac{e^{z_{2,k}}}{{\sum }_{j=1}^K{e}^{z_{2,j}}}$$

and for one-hot target $y$ the cross-entropy loss is

$$\mathcal{L}(\hat{y},y)=-\sum _{k=1}^K{y}_k\log {\hat{y}}_k$$

If the correct class is $c$, this reduces to

$$\mathcal{L}=-\log {\hat{y}}_c$$

Backpropagation Derivation

The output-layer derivative is the standard identity:

$$\frac{\partial \mathcal{L}}{\partial z_2}=\hat{y}-y$$

This falls out of differentiating the softmax and cross-entropy together. The rest is the chain rule.

The forward pass builds the computation left to right; the backward pass walks the same graph in reverse:

flowchart LR
    x["x"] --> z1["z1 = W1 x + b1"]
    z1 --> a1["a1 = phi(z1)"]
    a1 --> z2["z2 = W2 a1 + b2"]
    z2 --> yhat["y-hat = softmax(z2)"]
    yhat --> L["loss"]
    L -. "dL/dz2 = y-hat - y" .-> z2
    z2 -. "dL/da1 = W2^T dz2" .-> a1
    a1 -. "dL/dz1 = da1 * phi'(z1)" .-> z1

The solid edges are the forward pass; the dotted edges are the gradients flowing back through the same nodes.

Backpropagation is local

Each node only needs two things: the gradient arriving from above and the values it cached during the forward pass. The parameter gradients below are outer products of exactly those two quantities — $\partial \mathcal{L}/\partial W_2=(\hat{y}-y)a_1^{\top }$ pairs the upstream gradient with the cached input. This locality is why the forward pass must store activations, and why autodiff frameworks can differentiate any composition of primitives without global analysis.

Output Layer

Because $z_2=W_2{a}_1+b_2$:

$$\frac{\partial \mathcal{L}}{\partial W_2}=\frac{\partial \mathcal{L}}{\partial z_2}\frac{\partial z_2}{\partial W_2}=(\hat{y}-y)a_1^{\top }$$

and

$$\frac{\partial \mathcal{L}}{\partial b_2}=\hat{y}-y$$

The hidden activation receives

$$\frac{\partial \mathcal{L}}{\partial a_1}=W_2^{\top }(\hat{y}-y)$$

Hidden Layer

Since $a_1=\varphi (z_1)$:

$$\frac{\partial \mathcal{L}}{\partial z_1}=\frac{\partial \mathcal{L}}{\partial a_1}\odot {\varphi }^{\prime }(z_1)$$

For ReLU,

$$\varphi (z)=\max (0,z),{\varphi }^{\prime }(z)=1[z>0]$$

so only active hidden units pass gradient.

Now

$$\frac{\partial \mathcal{L}}{\partial W_1}=\frac{\partial \mathcal{L}}{\partial z_1}{x}^{\top },\frac{\partial \mathcal{L}}{\partial b_1}=\frac{\partial \mathcal{L}}{\partial z_1}$$

This is the whole backpropagation pattern in miniature. Every later architecture is the same story, only with more structured Jacobians.

Batch Form

For batch matrix $X\in {\mathbb{R}}^{B\times d_{in}}$:

$$\begin{array}{rl}{Z}_1& =X{W}_1^{\top }+b_1\\ A_1& =\varphi (Z_1)\\ Z_2& =A_1{W}_2^{\top }+b_2\\ \hat{Y}& =\text{softmax}(Z_2)\end{array}$$

In practice the loss is averaged over the batch.

Output Heads and Loss Families

The softmax-cross-entropy pair above is one instance of a general pattern: match the output head to the loss, and the gradient at the logits collapses to prediction minus target (Deep Learning ch. 6.2). The three standard pairings:

Task	Head	Loss	$\partial \mathcal{L}/\partial z$
Regression	linear, $\hat{y}=z$	MSE $\frac{1}{2}(\hat{y}-y{)}^2$	$z-y$
Binary classification	sigmoid, $\hat{y}=\sigma (z)$	BCE $-y\log \hat{y}-(1-y)\log (1-\hat{y})$	$\sigma (z)-y$
Multiclass	softmax	cross-entropy $-{\sum }_k{y}_k\log {\hat{y}}_k$	$\hat{y}-y$

Each row is the negative log-likelihood of a distribution family (Gaussian, Bernoulli, multinomial) under its canonical parameterization, which is why the clean form is not a coincidence.

The binary case is the promised second worked derivative. With $\hat{y}=\sigma (z)$ and ${\sigma }^{\prime }(z)=\sigma (z)(1-\sigma (z))$:

$$\frac{\partial \mathcal{L}}{\partial z}=\left(-\frac{y}{\hat{y}}+\frac{1-y}{1-\hat{y}}\right)\hat{y}(1-\hat{y})=-y(1-\hat{y})+(1-y)\hat{y}=\sigma (z)-y$$

The $\hat{y}(1-\hat{y})$ factor from the sigmoid exactly cancels the denominators from the log loss. This cancellation is also the numerical-stability argument for computing loss from logits (CrossEntropyLoss, binary_cross_entropy_with_logits) rather than from probabilities: the fused form never materializes a $\log$ of a saturated sigmoid or softmax. Pairing MSE with a sigmoid head, by contrast, leaves a ${\sigma }^{\prime }(z)$ factor in the gradient that vanishes whenever the unit saturates, which is a classic slow-training bug.

Optimizer Updates

The step method below implements plain SGD: $\theta \leftarrow \theta -\alpha g$. Two upgrades cover most practice; full treatment, including conditioning theory and a measured comparison, is in Numerical Optimization for Machine Learning.

Momentum accumulates a velocity so that consistent gradient directions compound and oscillating ones cancel:

$$v\leftarrow \mu v-\alpha g,\theta \leftarrow \theta +v$$

with $\mu$ typically 0.9 (Sutskever et al. 2013 schedule it from 0.5 up to 0.99). Adam keeps per-parameter moving averages of the gradient and its square, corrects their zero-initialization bias, and scales each coordinate’s step:

$$\begin{array}{rlrl}m& \leftarrow {\beta }_{1}m+(1-{\beta }_1)g,& \hat{m}& =m/(1-{\beta }_1^t)\\ v& \leftarrow {\beta }_{2}v+(1-{\beta }_2)g^2,& \hat{v}& =v/(1-{\beta }_2^t)\\ \theta & \leftarrow \theta -\alpha \hat{m}/(\sqrt{\hat{v}}+ϵ)& & \end{array}$$

with paper defaults $\alpha =10^{-3}$, ${\beta }_1=0.9$, ${\beta }_2=0.999$, $ϵ=10^{-8}$. As an implementation sketch on the MLP class below, momentum is three lines:

def step_momentum(self, grads, state, lr=1e-2, mu=0.9):
    for name in ("W1", "b1", "W2", "b2"):
        state[name] = mu * state.get(name, 0.0) - lr * grads[name]
        setattr(self, name, getattr(self, name) + state[name])

Schedules and Gradient Clipping

Two knobs sit outside the optimizer proper. Learning-rate schedules decay $\alpha$ over training — step decay (multiply by 0.1 every $N$ epochs), cosine annealing ${\alpha }_t=\frac{{\alpha }_0}{2}(1+\cos (\pi t/T))$, and a linear warmup over the first few hundred or thousand steps, which matters most with Adam because its second-moment estimate is unreliable early (DL ch. 8.5). Gradient clipping bounds the update when the loss surface produces a rare enormous gradient, as in recurrent nets (DL ch. 10.11.1); clip-by-norm rescales the whole gradient vector when it exceeds a threshold $\tau$:

$$g\leftarrow \{\begin{array}{ll}g& \Vert g\Vert \le \tau \\ \tau g/\Vert g\Vert & \text{otherwise}\end{array}$$

which preserves direction, unlike elementwise clip-by-value.

Gradient Checking

Before trusting a hand-written backward pass, compare it against a centered finite difference:

$$\frac{\partial \mathcal{L}}{\partial {\theta }_i}\approx \frac{\mathcal{L}({\theta }_i+h)-\mathcal{L}({\theta }_i-h)}{2h}$$

which has $O(h^2)$ error versus $O(h)$ for the one-sided version. The comparison metric is relative error $|g_a-g_n|/\max (|g_a|,|g_n|)$; per the CS231n conventions, below $10^{-7}$ is excellent for smooth networks, up to $10^{-4}$ is acceptable when ReLU kinks are involved, and above $10^{-2}$ means a bug. Use float64 (float32 cancellation error alone can reach $10^{-2}$), $h\approx 10^{-5}$, a handful of examples, and turn off dropout and other stochastic parts while checking.

Checking one entry of each weight matrix of the NumPy MLP below (run in the repo venv, float64, $h=10^{-5}$):

i, j, h = 1, 2, 1e-5
orig = model.W1[i, j]
model.W1[i, j] = orig + h; lp, _ = model.loss_and_grads(X, y)
model.W1[i, j] = orig - h; lm, _ = model.loss_and_grads(X, y)
model.W1[i, j] = orig
numeric = (lp - lm) / (2 * h)

Measured: W1[1,2] analytic $+0.35312938$ vs numeric $+0.35312938$, relative error $4.2\times 10^{-12}$; W2[1,2] relative error $1.7\times 10^{-10}$. Both are far inside the “excellent” band, which is the expected outcome for this loss because softmax-cross-entropy is smooth in the parameters even though ReLU has a kink in the inputs — the check would only brush the kink if a perturbation flipped a unit’s sign.

Why Initialization Matters

Very deep networks fail easily if activations or gradients change scale too aggressively across layers.

He et al. derive an initialization for rectifier networks that keeps the forward variance roughly stable:

$$\mathrm{Var}[W_{ij}]=\frac{2}{n_{in}}$$

so one common choice is

$$W_{ij}\sim \mathcal{N}\left(0,\frac{2}{n_{in}}\right)$$

That is the standard He initialization. The factor of $2$ appears because ReLU zeroes about half the mass.

The same paper also proposes PReLU:

$$f(y_i)=\{\begin{array}{ll}{y}_i& y_i>0\\ a_i{y}_i& y_i\le 0\end{array}$$

with learned negative slope $a_i$. The paper reports 4.94% top-5 test error on ImageNet 2012, a 26% relative improvement over GoogLeNet’s 6.66%.

Batch Normalization

Ioffe and Szegedy normalize each activation dimension over the current mini-batch:

$${\hat{x}}^{(k)}=\frac{x^{(k)}-{\mathbb{E}}_{\mathcal{B}}[x^{(k)}]}{\sqrt{{\mathrm{Var}}_{\mathcal{B}}[x^{(k)}]+ϵ}}$$

then restore learnable scale and shift:

$$y^{(k)}={\gamma }^{(k)}{\hat{x}}^{(k)}+{\beta }^{(k)}$$

The point is not to clamp everything permanently to zero mean and unit variance. The point is to stabilize optimization while still letting the model learn whatever affine reparameterization it needs.

The BatchNorm paper reports the same accuracy in 14 times fewer training steps on a strong ImageNet model, and an ensemble reaching 4.9% top-5 validation error.

Dropout

Dropout randomly removes units during training. If $m$ is a Bernoulli mask:

$$\tilde{h}=m\odot h$$

The paper’s framing is useful. Training samples an exponential family of “thinned” subnetworks, and test-time inference approximates their average with one full network.

Srivastava et al. describe dropout as a way to prevent units from co-adapting too much. In several experiments they found dropping 20% of input units and 50% of hidden units worked well.

Choosing Initialization and Normalization

The pieces above interact, and the common configurations are worth a table:

Choice	Formula / rule	When
Xavier/Glorot init	$\mathrm{Var}[W]=\frac{2}{n_{in}+n_{out}}$	tanh or sigmoid activations
He init	$\mathrm{Var}[W]=\frac{2}{n_{in}}$	ReLU-family activations
Batch norm	normalize per-dimension over the batch	large batches; convolutional nets
Layer norm	normalize per-example over features	small/variable batches; transformers
L2 / weight decay	add $\frac{\lambda }{2}\Vert \theta {\Vert }^2$ (or decoupled decay)	default regularizer
Dropout	Bernoulli mask, $p\approx 0.5$ hidden / $0.2$ input	large fully connected layers

The interactions: normalization layers make initialization scale much less critical, because activations get renormalized each layer regardless of what the weights did — with batch norm, a wrong init costs early training speed rather than trainability. Weight decay and He/Xavier initialization pull in compatible directions (both keep weights in a moderate range), but decay interacts with adaptive optimizers in a way that matters; see the decoupled weight decay discussion in Numerical Optimization. Dropout raises activation variance during training, which is why implementations use inverted dropout (rescale by $1/(1-p)$ at train time) so that test-time inference needs no correction.

NumPy Implementation

This version does the forward pass, gradient calculation, and SGD update explicitly.

import numpy as np
 
def relu(x):
    return np.maximum(x, 0.0)
 
def relu_grad(x):
    return (x > 0).astype(x.dtype)
 
def softmax(logits):
    shifted = logits - logits.max(axis=1, keepdims=True)
    exp = np.exp(shifted)
    return exp / exp.sum(axis=1, keepdims=True)
 
class MLP:
    def __init__(self, d_in, d_hidden, d_out, seed=0):
        rng = np.random.default_rng(seed)
        self.W1 = rng.normal(0.0, np.sqrt(2 / d_in), size=(d_hidden, d_in))
        self.b1 = np.zeros(d_hidden)
        self.W2 = rng.normal(0.0, np.sqrt(2 / d_hidden), size=(d_out, d_hidden))
        self.b2 = np.zeros(d_out)
 
    def forward(self, X):
        z1 = X @ self.W1.T + self.b1
        a1 = relu(z1)
        z2 = a1 @ self.W2.T + self.b2
        probs = softmax(z2)
        cache = (X, z1, a1, probs)
        return probs, cache
 
    def loss_and_grads(self, X, y):
        probs, (X, z1, a1, probs) = self.forward(X)
        B = X.shape[0]
 
        one_hot = np.zeros_like(probs)
        one_hot[np.arange(B), y] = 1.0
        loss = -np.log(probs[np.arange(B), y] + 1e-12).mean()
 
        dz2 = (probs - one_hot) / B
        dW2 = dz2.T @ a1
        db2 = dz2.sum(axis=0)
 
        da1 = dz2 @ self.W2
        dz1 = da1 * relu_grad(z1)
        dW1 = dz1.T @ X
        db1 = dz1.sum(axis=0)
 
        grads = {"W1": dW1, "b1": db1, "W2": dW2, "b2": db2}
        return loss, grads
 
    def step(self, grads, lr=1e-2):
        self.W1 -= lr * grads["W1"]
        self.b1 -= lr * grads["b1"]
        self.W2 -= lr * grads["W2"]
        self.b2 -= lr * grads["b2"]

PyTorch Implementation

import torch
import torch.nn as nn
 
class TorchMLP(nn.Module):
    def __init__(self, d_in: int, d_hidden: int, d_out: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_in, d_hidden),
            nn.ReLU(),
            nn.Linear(d_hidden, d_out),
        )
 
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)
 
model = TorchMLP(128, 256, 10)
logits = model(torch.randn(32, 128))
loss = nn.CrossEntropyLoss()(logits, torch.randint(0, 10, (32,)))
loss.backward()

This is shorter because autograd is doing exactly the backpropagation derivation from above.

Executable Experiments

The first notebook trains the NumPy model on MNIST and exposes the loss curve, predictions, and failure cases discussed above.

Jupyter Notebook

mlp-from-scratch-on-mnist.ipynb

Goal

Build a small multilayer perceptron from scratch with NumPy, train it on the real MNIST dataset, and inspect where the implementation earns its accuracy. The main point is not raw leaderboard performance. The point is to expose each tensor in the forward and backward pass.

Why This Dataset

MNIST is small enough to train on a laptop and still rich enough to show the shape of supervised learning. The input is a (28 \times 28) image flattened into (784) features. The target is one of (10) classes.

We will fit

$$x\in {\mathbb{R}}^{784}\to h=\text{ReLU}(x{W}_1+b_1)\to \hat{y}=\text{softmax}(h{W}_2+b_2).$$

The loss for a minibatch of size (B) is

$$\mathcal{L}=-\frac{1}{B}\sum _{i=1}^B\log p_{i,y_i}.$$

In [1]:

## Uncomment this if you are running in a fresh environment.
## %pip install -q datasets numpy matplotlib
 
import math
import random
from dataclasses import dataclass
 
import matplotlib.pyplot as plt
import numpy as np
from datasets import load_dataset
 
np.random.seed(7)
random.seed(7)

In [2]:

ds = load_dataset("ylecun/mnist")
 
x_train = np.stack([np.array(example["image"], dtype=np.float32).reshape(-1) for example in ds["train"]]) / 255.0
y_train = np.array(ds["train"]["label"], dtype=np.int64)
x_test = np.stack([np.array(example["image"], dtype=np.float32).reshape(-1) for example in ds["test"]]) / 255.0
y_test = np.array(ds["test"]["label"], dtype=np.int64)
 
x_train = x_train[:20000]
y_train = y_train[:20000]
 
print(x_train.shape, y_train.shape, x_test.shape, y_test.shape)

Out[2]:

(20000, 784) (20000,) (10000, 784) (10000,)

NumPy Model

The hidden layer is

$$H=\max (0,X{W}_1+b_1)$$

and the logits are

$$Z=H{W}_2+b_2.$$

After the softmax, the derivative of cross-entropy with respect to the logits is

$$\frac{\partial \mathcal{L}}{\partial Z}=\frac{P-Y}{B},$$

where (P) is the predicted probability matrix and (Y) is the one-hot target matrix.

In [3]:

def one_hot(y: np.ndarray, num_classes: int) -> np.ndarray:
    out = np.zeros((len(y), num_classes), dtype=np.float32)
    out[np.arange(len(y)), y] = 1.0
    return out
 
 
def relu(x: np.ndarray) -> np.ndarray:
    return np.maximum(x, 0.0)
 
 
def relu_grad(x: np.ndarray) -> np.ndarray:
    return (x > 0).astype(np.float32)
 
 
def softmax(logits: np.ndarray) -> np.ndarray:
    shifted = logits - logits.max(axis=1, keepdims=True)
    exp = np.exp(shifted)
    return exp / exp.sum(axis=1, keepdims=True)
 
 
def cross_entropy(probs: np.ndarray, y: np.ndarray) -> float:
    eps = 1e-9
    return float(-np.log(probs[np.arange(len(y)), y] + eps).mean())
 
 
@dataclass
class MLP:
    hidden_dim: int = 256
    input_dim: int = 784
    output_dim: int = 10
 
    def __post_init__(self) -> None:
        scale1 = math.sqrt(2.0 / self.input_dim)
        scale2 = math.sqrt(2.0 / self.hidden_dim)
        self.W1 = np.random.randn(self.input_dim, self.hidden_dim).astype(np.float32) * scale1
        self.b1 = np.zeros((1, self.hidden_dim), dtype=np.float32)
        self.W2 = np.random.randn(self.hidden_dim, self.output_dim).astype(np.float32) * scale2
        self.b2 = np.zeros((1, self.output_dim), dtype=np.float32)
 
    def forward(self, x: np.ndarray) -> tuple[np.ndarray, dict[str, np.ndarray]]:
        z1 = x @ self.W1 + self.b1
        h1 = relu(z1)
        z2 = h1 @ self.W2 + self.b2
        probs = softmax(z2)
        cache = {"x": x, "z1": z1, "h1": h1, "z2": z2, "probs": probs}
        return probs, cache
 
    def backward(self, cache: dict[str, np.ndarray], y: np.ndarray) -> dict[str, np.ndarray]:
        x = cache["x"]
        z1 = cache["z1"]
        h1 = cache["h1"]
        probs = cache["probs"].copy()
        batch_size = len(y)
 
        probs[np.arange(batch_size), y] -= 1.0
        probs /= batch_size
 
        dW2 = h1.T @ probs
        db2 = probs.sum(axis=0, keepdims=True)
        dh1 = probs @ self.W2.T
        dz1 = dh1 * relu_grad(z1)
        dW1 = x.T @ dz1
        db1 = dz1.sum(axis=0, keepdims=True)
        return {"W1": dW1, "b1": db1, "W2": dW2, "b2": db2}
 
    def step(self, grads: dict[str, np.ndarray], lr: float) -> None:
        self.W1 -= lr * grads["W1"]
        self.b1 -= lr * grads["b1"]
        self.W2 -= lr * grads["W2"]
        self.b2 -= lr * grads["b2"]

In [4]:

def accuracy(model: MLP, x: np.ndarray, y: np.ndarray, batch_size: int = 512) -> float:
    preds = []
    for start in range(0, len(x), batch_size):
        xb = x[start : start + batch_size]
        probs, _ = model.forward(xb)
        preds.append(probs.argmax(axis=1))
    pred = np.concatenate(preds)
    return float((pred == y).mean())
 
 
def iterate_minibatches(x: np.ndarray, y: np.ndarray, batch_size: int):
    idx = np.random.permutation(len(x))
    for start in range(0, len(x), batch_size):
        batch_idx = idx[start : start + batch_size]
        yield x[batch_idx], y[batch_idx]
 
 
model = MLP(hidden_dim=256)
history = {"train_loss": [], "train_acc": [], "test_acc": []}
 
epochs = 12
lr = 0.08
batch_size = 256
 
for epoch in range(epochs):
    losses = []
    for xb, yb in iterate_minibatches(x_train, y_train, batch_size):
        probs, cache = model.forward(xb)
        loss = cross_entropy(probs, yb)
        grads = model.backward(cache, yb)
        model.step(grads, lr=lr)
        losses.append(loss)
 
    train_acc = accuracy(model, x_train[:4000], y_train[:4000])
    test_acc = accuracy(model, x_test, y_test)
    history["train_loss"].append(float(np.mean(losses)))
    history["train_acc"].append(train_acc)
    history["test_acc"].append(test_acc)
    print(f"epoch={epoch:02d} loss={history['train_loss'][-1]:.4f} train_acc={train_acc:.4f} test_acc={test_acc:.4f}")

Out[4]:

epoch=00 loss=1.0016 train_acc=0.8632 test_acc=0.8637
epoch=01 loss=0.4673 train_acc=0.8858 test_acc=0.8780
epoch=02 loss=0.3784 train_acc=0.9095 test_acc=0.9027
epoch=03 loss=0.3385 train_acc=0.9193 test_acc=0.9061

epoch=04 loss=0.3081 train_acc=0.9283 test_acc=0.9155

epoch=05 loss=0.2881 train_acc=0.9300 test_acc=0.9202
epoch=06 loss=0.2729 train_acc=0.9353 test_acc=0.9221
epoch=07 loss=0.2587 train_acc=0.9355 test_acc=0.9218
epoch=08 loss=0.2465 train_acc=0.9420 test_acc=0.9274

epoch=09 loss=0.2387 train_acc=0.9435 test_acc=0.9268

epoch=10 loss=0.2259 train_acc=0.9435 test_acc=0.9330
epoch=11 loss=0.2192 train_acc=0.9477 test_acc=0.9306

In [5]:

fig, axes = plt.subplots(1, 2, figsize=(12, 4))
axes[0].plot(history["train_loss"], marker="o")
axes[0].set_title("Training loss")
axes[0].set_xlabel("epoch")
axes[0].set_ylabel("cross entropy")
 
axes[1].plot(history["train_acc"], marker="o", label="train")
axes[1].plot(history["test_acc"], marker="o", label="test")
axes[1].set_title("Accuracy")
axes[1].set_xlabel("epoch")
axes[1].legend()
plt.show()

Out[5]:

<Figure size 1200x400 with 2 Axes>

In [6]:

sample_idx = np.random.choice(len(x_test), size=12, replace=False)
fig, axes = plt.subplots(3, 4, figsize=(10, 8))
probs, _ = model.forward(x_test[sample_idx])
preds = probs.argmax(axis=1)
 
for ax, idx, pred in zip(axes.flat, sample_idx, preds):
    ax.imshow(x_test[idx].reshape(28, 28), cmap="gray")
    ax.set_title(f"pred={pred} true={y_test[idx]}")
    ax.axis("off")
 
plt.tight_layout()
plt.show()

Out[6]:

<Figure size 1000x800 with 12 Axes>

What To Extend

A few natural next steps:

1. Replace plain SGD with momentum or Adam and compare convergence.
2. Add another hidden layer and watch how the gradient norms change.
3. Swap the full-batch NumPy implementation for a Torch version and compare ergonomics.

The second isolates softmax regression so the symbolic gradient, finite differences, and PyTorch autograd can be compared on the same real minibatch.

Jupyter Notebook

checking-manual-gradients-against-autodiff-on-mnist.ipynb

Goal

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, derive the gradient of softmax regression by hand, and compare it against torch.autograd.

For logits (Z = XW + b), probabilities (P = \text{softmax}(Z)), and one-hot targets (Y), the cross-entropy derivative is

$$\frac{\partial \mathcal{L}}{\partial Z}=\frac{P-Y}{B}.$$

The weight gradient follows immediately:

$$\frac{\partial \mathcal{L}}{\partial W}=X^{\top }\frac{P-Y}{B}.$$

In [1]:

## %pip install -q datasets torch numpy matplotlib
 
import numpy as np
import torch
import torch.nn.functional as F
from datasets import load_dataset

In [2]:

ds = load_dataset("ylecun/mnist")
train_subset = ds["train"].select(range(1024))
test_subset = ds["test"].select(range(1024))
 
x = np.stack([np.array(example["image"], dtype=np.float32).reshape(-1) for example in train_subset]) / 255.0
y = np.array(train_subset["label"], dtype=np.int64)
 
batch_x = x[:256]
batch_y = y[:256]
 
print(batch_x.shape, batch_y.shape)

Out[2]:

(256, 784) (256,)

In [3]:

def softmax_numpy(logits: np.ndarray) -> np.ndarray:
    shifted = logits - logits.max(axis=1, keepdims=True)
    exp = np.exp(shifted)
    return exp / exp.sum(axis=1, keepdims=True)
 
 
def manual_gradients(xb: np.ndarray, yb: np.ndarray, W: np.ndarray, b: np.ndarray):
    logits = xb @ W + b
    probs = softmax_numpy(logits)
    one_hot = np.zeros_like(probs)
    one_hot[np.arange(len(yb)), yb] = 1.0
    dz = (probs - one_hot) / len(yb)
    dW = xb.T @ dz
    db = dz.sum(axis=0, keepdims=True)
    loss = -np.log(probs[np.arange(len(yb)), yb] + 1e-9).mean()
    return loss, dW, db

In [4]:

torch.manual_seed(7)
 
W_np = np.random.randn(784, 10).astype(np.float32) * 0.01
b_np = np.zeros((1, 10), dtype=np.float32)
 
manual_loss, manual_dW, manual_db = manual_gradients(batch_x, batch_y, W_np, b_np)
 
xb_t = torch.tensor(batch_x, dtype=torch.float32)
yb_t = torch.tensor(batch_y, dtype=torch.long)
W_t = torch.tensor(W_np, dtype=torch.float32, requires_grad=True)
b_t = torch.tensor(b_np, dtype=torch.float32, requires_grad=True)
 
logits_t = xb_t @ W_t + b_t
loss_t = F.cross_entropy(logits_t, yb_t)
loss_t.backward()
 
max_abs_dW = np.abs(manual_dW - W_t.grad.detach().numpy()).max()
max_abs_db = np.abs(manual_db - b_t.grad.detach().numpy()).max()
 
print("manual loss:", manual_loss)
print("autodiff loss:", float(loss_t))
print("max |dW_manual - dW_autodiff| =", max_abs_dW)
print("max |db_manual - db_autodiff| =", max_abs_db)

Out[4]:

manual loss: 2.3231063
autodiff loss: 2.323106527328491
max |dW_manual - dW_autodiff| = 1.4901161e-08
max |db_manual - db_autodiff| = 1.8626451e-08

Finite Difference Sanity Check

Autodiff and the symbolic derivation can still agree on the same bug. A cheap extra check is finite differences on a few random coordinates:

$$\frac{\partial \mathcal{L}}{\partial {\theta }_j}\approx \frac{\mathcal{L}({\theta }_j+\epsilon )-\mathcal{L}({\theta }_j-\epsilon )}{2\epsilon }.$$

In [5]:

def loss_only(xb: np.ndarray, yb: np.ndarray, W: np.ndarray, b: np.ndarray) -> float:
    logits = xb @ W + b
    probs = softmax_numpy(logits)
    return float(-np.log(probs[np.arange(len(yb)), yb] + 1e-9).mean())
 
 
eps = 1e-4
checks = [(0, 0), (123, 3), (511, 7)]
 
for i, j in checks:
    W_pos = W_np.copy()
    W_neg = W_np.copy()
    W_pos[i, j] += eps
    W_neg[i, j] -= eps
    fd = (loss_only(batch_x, batch_y, W_pos, b_np) - loss_only(batch_x, batch_y, W_neg, b_np)) / (2 * eps)
    print((i, j), "finite-diff =", fd, "manual =", manual_dW[i, j])

Out[5]:

(0, 0) finite-diff = 0.0 manual = 0.0
(123, 3) finite-diff = -0.00476837158203125 manual = -0.004906716
(511, 7) finite-diff = 0.015497207641601562 manual = 0.01614688

In [6]:

class SoftmaxRegressor(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(784, 10)
 
    def forward(self, x):
        return self.linear(x)
 
 
model = SoftmaxRegressor()
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
 
x_train = torch.tensor(x[:1024], dtype=torch.float32)
y_train = torch.tensor(y[:1024], dtype=torch.long)
x_test = torch.tensor(
    np.stack([np.array(example["image"], dtype=np.float32).reshape(-1) for example in test_subset]) / 255.0,
    dtype=torch.float32,
)
y_test = torch.tensor(test_subset["label"], dtype=torch.long)
 
history = []
batch_size = 256
 
for epoch in range(10):
    perm = torch.randperm(len(x_train))
    for start in range(0, len(x_train), batch_size):
        idx = perm[start : start + batch_size]
        logits = model(x_train[idx])
        loss = F.cross_entropy(logits, y_train[idx])
        opt.zero_grad()
        loss.backward()
        opt.step()
 
    with torch.no_grad():
        train_acc = (model(x_train).argmax(dim=1) == y_train).float().mean().item()
        test_acc = (model(x_test).argmax(dim=1) == y_test).float().mean().item()
        history.append((train_acc, test_acc))
        print(f"epoch={epoch:02d} train_acc={train_acc:.4f} test_acc={test_acc:.4f}")

Out[6]:

epoch=00 train_acc=0.3193 test_acc=0.2715
epoch=01 train_acc=0.5527 test_acc=0.4365
epoch=02 train_acc=0.6621 test_acc=0.5186
epoch=03 train_acc=0.7275 test_acc=0.6084
epoch=04 train_acc=0.7646 test_acc=0.6396
epoch=05 train_acc=0.7842 test_acc=0.6758
epoch=06 train_acc=0.8018 test_acc=0.6895
epoch=07 train_acc=0.8115 test_acc=0.7002
epoch=08 train_acc=0.8125 test_acc=0.7090
epoch=09 train_acc=0.8164 test_acc=0.7168

Takeaway

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.

What to Carry Forward

• Backpropagation is just repeated local application of the chain rule.
• Match the output head to the loss; the logit gradient becomes prediction minus target.
• Gradient-check any hand-written backward pass before trusting it.
• Initialization is part of the model, not a clerical detail.
• Normalization layers often change optimization more than small architecture tweaks do.
• Dropout is an ensemble-style regularizer implemented inside one training loop.

Related Notes

• Modeling, Architecture, and Data
• Convolutional Neural Networks
• Decoder-Only Transformers
• Recurrent Neural Networks
• Numerical Optimization for Machine Learning
• Matrix Calculus for Machine Learning

Sources

• Rumelhart, Hinton, and Williams (1986), Learning Representations by Back-Propagating Errors
• He et al. (2015), Delving Deep into Rectifiers
• Ioffe and Szegedy (2015), Batch Normalization
• Srivastava et al. (2014), Dropout
• Goodfellow, Bengio, and Courville, Deep Learning, chapters 6, 8, and 10
• Kingma and Ba (2015), Adam: A Method for Stochastic Optimization
• Sutskever et al. (2013), On the importance of initialization and momentum in deep learning
• CS231n, Neural Networks Part 3: Gradient checks