A staged plan for reimplementing a paper

Setup#

pip install scikit-learn

Verified with scikit-learn 1.7.2 and numpy 1.26.4 on CPU. The whole script runs in a few seconds and downloads nothing — load_digits ships inside scikit-learn.

The worked target#

We continue with Kingma and Ba (2015). Section 6.1 of the paper trains multi-class logistic regression on MNIST with Adam. Our scaled-down target: train the same model with a hand-written Adam and land near a trusted implementation of the same model.

Be clear about the scaling down. load_digits is 1,797 images of 8×8 pixels — a different, far easier dataset than MNIST's 70,000 images of 28×28. We are reproducing the shape of the experiment on a laptop, not its numbers. Saying so out loud is part of doing this honestly.

Stages 3, 4 and 5 in one script#

import numpy as np
from sklearn.datasets import load_digits
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split

def softmax_xent(W, b, X, Y):
    """Loss and its gradients for multi-class logistic regression."""
    z = X @ W + b
    z = z - z.max(axis=1, keepdims=True)               # stable log-sum-exp
    logp = z - np.log(np.exp(z).sum(axis=1, keepdims=True))
    loss = -(Y * logp).mean(axis=0).sum()
    gz = (np.exp(logp) - Y) / len(X)
    return loss, X.T @ gz, gz.sum(axis=0)

def fit(X, Y, steps, alpha=0.01, seed=0):
    rng = np.random.default_rng(seed)
    W = rng.normal(0, 0.01, (X.shape[1], Y.shape[1]))
    b = np.zeros(Y.shape[1])
    mW, vW = np.zeros_like(W), np.zeros_like(W)
    mb, vb = np.zeros_like(b), np.zeros_like(b)
    for t in range(1, steps + 1):
        loss, gW, gb = softmax_xent(W, b, X, Y)
        for p, g, m, v in ((W, gW, mW, vW), (b, gb, mb, vb)):
            m *= 0.9                                   # m and v persist across steps
            m += 0.1 * g
            v *= 0.999
            v += 0.001 * g * g
            p -= alpha * (m / (1 - 0.9**t)) / (np.sqrt(v / (1 - 0.999**t)) + 1e-8)
    return W, b, loss

d = load_digits()
X, Y = d.data / 16.0, np.eye(10)[d.target]
Xtr, Xte, Ytr, Yte = train_test_split(X, Y, test_size=0.3, random_state=0)

# stage 3 -- can it memorise eight examples? if not, the update is broken.
_, _, small = fit(Xtr[:8], Ytr[:8], steps=500)
print(f"stage 3  loss on 8 examples after 500 steps: {small:.6f}")

# stage 4 -- the real (small) dataset
W, b, loss = fit(Xtr, Ytr, steps=2000)
acc = (np.argmax(Xte @ W + b, axis=1) == np.argmax(Yte, axis=1)).mean()
print(f"stage 4  training loss {loss:.4f}   test accuracy {acc:.4f}")

# stage 5 -- against a trusted implementation of the same model
ref = LogisticRegression(max_iter=5000).fit(Xtr, np.argmax(Ytr, axis=1))
print(f"stage 5  sklearn LogisticRegression test accuracy "
      f"{ref.score(Xte, np.argmax(Yte, axis=1)):.4f}")

stage 3  loss on 8 examples after 500 steps: 0.003988
stage 4  training loss 0.0167   test accuracy 0.9574
stage 5  sklearn LogisticRegression test accuracy 0.9611

Read the three lines as three verdicts.

Stage 3 passed. Loss near zero on eight examples means gradients flow, labels line up with predictions, and the optimiser moves parameters in the right direction. Everything downstream is now worth debugging.

Stage 4 produced a number, which is meaningless on its own.

Stage 5 gave it meaning. 95.74% against 96.11% from a mature, well-tested implementation of the same model. A gap of 0.37 points on the same data is the neighbourhood you want — close enough that the method is right, far enough that something differs. Chasing that remainder is the next lesson.

The walkthrough#

Stage 3 costs six seconds and rules out half of all bugs. A model that cannot memorise eight examples has a plumbing fault, not a tuning problem. This is overfitting one batch, applied before there is anything to overfit.

The optimiser is the transcribed Adam from the previous lesson, now updating two parameter arrays rather than one. m and v are mutated in place inside the tuple loop, so their state survives between steps. That in-place update is easy to break during a refactor — assert m.any() after the first step if you ever suspect it.

softmax_xent subtracts the row maximum before exponentiating. Straight transcription of the written formula overflows on large scores; this rewrite is algebraically identical and stable.

The reference is the same model, not the same paper. scikit-learn's LogisticRegression uses a different solver and adds a small L2 penalty by default. It is a sanity anchor, not a ground truth. Comparing against a different model would tell you nothing.

The reported loss comes from the last forward pass, before the final update. At 2,000 steps that distinction is invisible; at 5 steps it is not.

Budget the project before starting it#

Before stage 1, write down three numbers: how long one full run takes, how many full runs you can afford, and what a full run costs in money. A paper that trained for 500 GPU-days cannot be reproduced end to end on a laptop, and knowing that on day one changes the project into something achievable — one ablation, one dataset, one figure.

Then scale down deliberately and write down what you changed: smaller dataset, fewer layers, shorter schedule. A scaled-down reproduction that states its scaling is a real result. An unstated one is a broken result.

Common mistakes#

Starting with the interesting part. The novel mechanism is the fun bit and the wrong bit to write first. Without evaluation, you cannot tell whether it is working.

No baseline in the same script. A number with nothing beside it cannot be interpreted. Keep a trusted reference and a dumb baseline running in every experiment, forever.

Fixing stage 4 by tuning. When the small run underperforms, the reflex is to raise the learning rate or add epochs. If stage 3 was skipped, you may be tuning around a broken gradient. Go back rather than forward.

Changing several things per run. Once results are close, change one thing at a time and record what happened. See ablation studies.

Not writing the target down. Without stage 0 in a file, the project ends when you get bored rather than when you are done.

Try it yourself#

Break stage 3 on purpose: shuffle Ytr so labels no longer match their images, and rerun. The loss on eight examples will still fall near zero — memorising noise is easy — while stage 4 collapses to roughly 10%. That contrast is exactly why stage 3 checks plumbing and stage 4 checks learning, and why you need both.

What to learn next#

• Why your reimplementation is three points worse — closing the last gap.
• Overfit one batch — stage 3 as a permanent habit.
• Ablation studies — the discipline that keeps late-stage debugging interpretable.