Demo: Gradient Descent on MNIST

We will train a linear softmax classifier on handwritten digits, compare training with held-out performance, and inspect what its class weights represent. The first cell loads and scales a manageable subset of MNIST. The first run downloads MNIST from OpenML. Later runs reuse its local cache.

import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import fetch_openml

plt.style.use("../figure_style.mplstyle")
rng = np.random.default_rng(10)

X_all, y_all = fetch_openml("mnist_784", version=1, return_X_y=True, as_frame=False, parser="auto")
y_all = y_all.astype(int)

# subsample for a demo that trains in seconds, not minutes
idx = rng.choice(len(X_all), size=6000, replace=False)
X_all, y_all = X_all[idx], y_all[idx]
X_all = X_all / 255.0                      # pixels to [0, 1]
X_all = np.column_stack([X_all, np.ones(len(X_all))]).T  # examples in columns, constant feature in last row

n_train = 5000
X_train, y_train = X_all[:, :n_train], y_all[:n_train]
X_val, y_val = X_all[:, n_train:], y_all[n_train:]
print(f"training on {n_train} digits, validating on {len(y_val)}")

training on 5000 digits, validating on 1000

Question 1: What Computation Implements the Softmax Gradient?

The reading derives the probability-minus-label residual for each class. For \(n\) examples with \(d\) features and \(k\) classes, let \(\mathbf X\in\mathbb R^{d\times n}\) contain the feature vectors as columns and \(\mathbf W\in\mathbb R^{k\times d}\) contain one class’s weights per row. The score matrix is \(\mathbf Z=\mathbf W\mathbf X\). Let \(\mathbf P,\mathbf Y\in\mathbb R^{k\times n}\) contain the predicted probabilities and one-hot labels as columns. Averaging the per-example weight gradients gives

\[ \nabla_{\mathbf W}\mathcal L =\frac1n(\mathbf P-\mathbf Y)\mathbf X^\top. \]

The next cell builds the one-hot labels, stable softmax, loss, and accuracy needed for that update.

k = 10  # digit classes
d = X_train.shape[0]

# Column i has a 1 in the row for its correct digit.
Y_train = np.zeros((k, len(y_train)))
Y_train[y_train, np.arange(len(y_train))] = 1
Y_val = np.zeros((k, len(y_val)))
Y_val[y_val, np.arange(len(y_val))] = 1

def softmax(Z):
    Z = Z - Z.max(axis=0, keepdims=True)   # for numerical stability without changing the output
    expZ = np.exp(Z)
    return expZ / expZ.sum(axis=0, keepdims=True)

def cross_entropy(P, Y):
    return -np.mean(np.sum(Y * np.log(P + 1e-12), axis=0))

def accuracy(P, y):
    return np.mean(P.argmax(axis=0) == y)

Answer. Each column of \(\mathbf P-\mathbf Y\) is predicted probability minus the one-hot label. An update from one example raises its true-class score and lowers its competing scores. Multiplication by \(\mathbf X^\top\) combines all examples’ contributions to each weight. Because weights are shared, the average update need not improve every example.

Question 2: How Do Training and Validation Behave?

Starting from zero weights, we record both losses at the same parameter values. The model updates only from the training examples.

W = np.zeros((k, d))
alpha = 0.5
n_steps = 300

steps = []
train_losses = []
val_losses = []
val_accs = []
for t in range(n_steps):
    P_train = softmax(W @ X_train)

    if t % 5 == 0:   # both scores measured at the same weights, before the step
        steps.append(t)
        train_losses.append(cross_entropy(P_train, Y_train))
        P_val = softmax(W @ X_val)
        val_losses.append(cross_entropy(P_val, Y_val))
        val_accs.append(accuracy(P_val, y_val))

    grad = (P_train - Y_train) @ X_train.T / n_train      # softmax minus one-hot, averaged
    W -= alpha * grad

Compare training and validation loss at the same weights.

fig, ax = plt.subplots(figsize=(7, 3))
ax.plot(steps, train_losses, linestyle="--", linewidth=1.6, label="Training Cross-Entropy")
ax.plot(steps, val_losses, linewidth=1.6, label="Validation Cross-Entropy")
ax.set_xlabel("Gradient Descent Step")
ax.set_ylabel("Cross-Entropy Loss")
ax.legend(frameon=False)
plt.show()

print(f"final validation accuracy: {accuracy(softmax(W @ X_val), y_val):.1%}")

final validation accuracy: 90.9%

Answer. In this run, both losses fall, while validation loss remains higher. These held-out images show how well the fitted model predicts examples it did not train on. The final validation accuracy shows that even a linear classifier captures much of the digit structure.

Question 3: What Did Each Class Weight Learn?

Each row of \(\mathbf W\) contains one weight per pixel for one digit. After dropping the constant-feature entry and reshaping, positive and negative evidence becomes visible as an image.

fig, axes = plt.subplots(2, 5, figsize=(9, 4.4))
scale = np.percentile(np.abs(W[:, :-1]), 99)
for digit, ax in enumerate(axes.flat):
    weight_image = W[digit, :-1].reshape(28, 28)          # drop the constant-feature weight
    shown = ax.imshow(weight_image, cmap="RdBu", vmin=-scale, vmax=scale)
    ax.set_xlabel(f"Digit {digit}")
    ax.set_xticks([])
    ax.set_yticks([])
fig.colorbar(shown, ax=axes, fraction=0.025, pad=0.02, label="Class Weight")
plt.show()

Answer. The templates emphasize characteristic strokes and suppress pixels associated with competing digits. One template per class is expressive enough to be useful, but it cannot adapt its features to different handwriting styles.

Question 4: Does a Larger Step Size Always Train Faster?

Change alpha_new and compare validation accuracy across optimization steps. A useful step should move quickly without repeatedly overshooting good weights.

alpha_new = 3.0    # change me!

W2 = np.zeros((k, d))
val_accs2 = []
for t in range(n_steps):
    P_train = softmax(W2 @ X_train)
    if t % 5 == 0:
        val_accs2.append(accuracy(softmax(W2 @ X_val), y_val))
    gradient = (P_train - Y_train) @ X_train.T / n_train
    W2 -= alpha_new * gradient

Compare both learning rates after the same number of updates.

fig, ax = plt.subplots(figsize=(7, 3))
ax.plot(steps, val_accs, linewidth=1.6, label=f"$\\alpha={alpha}$")
ax.plot(steps, val_accs2, linestyle="--", linewidth=1.6, label=f"$\\alpha={alpha_new}$")
ax.set_xlabel("Gradient Descent Step")
ax.set_ylabel("Validation Accuracy")
ax.legend(frameon=False)
plt.show()

Answer. Increasing the learning rate can speed early progress, but an overly large value makes optimization unstable. The update direction and its step size solve different parts of the optimization problem.