A Module is three verbs — forward, backward, parameters — and a training flag. Sequential chains modules into networks. Loss objects compute the blame; optimizers spend it; a DataLoader feeds the loop. Build all of it in ~500 lines and PyTorch stops being magic: you will know exactly what every line of a training script is doing.
Ten lessons of pieces — a Value engine here, a training loop there. A framework makes them composable by agreeing that every layer implements the same interface: forward() computes, backward() passes gradients, parameters() exposes trainable scalars, and a training flag switches stochastic behaviour. Nothing about Linear, ReLU or Dropout is special beyond keeping that promise.
forward() · backward() · parameters() · train()/eval()The container is itself a Module — the composite pattern. Its forward pass feeds data through children left to right and caches along the way; its backward pass walks the same children right to left. Because parameters() concatenates every child's list, one optimizer can update a whole network it knows nothing about — 465 scalars for the source's 2→16→16→8→1 circle model.
2→16→16→8→1 = 48 + 272 + 136 + 9 = 465A DataLoader slices the dataset into batches and shuffles each epoch; a loss object computes the scalar and its gradient; an optimizer walks the flat parameter list. The five calls — zero_grad, forward, loss, backward, step — are the whole algorithm. Everything else in a framework is packaging around those lines.
zero_grad → forward → loss → backward → stepA framework is a contract plus a loop: every layer speaks the same three verbs, a container can stack anything that speaks them into a network, and training is just calling those verbs — forward, loss, backward, step — in the right order, on batches an iterator hands you.
By the end you will be able to: count a network’s parameters by hand and say what each one is for; walk a batch through a forward pass and name every shape; explain what each module caches and why backward needs it; step SGD and Adam by hand on a two-weight layer; explain what zero_grad does and why it is a separate call; say what train() and eval() change; and map every class in your framework to its PyTorch counterpart — nn.Module, nn.Sequential, nn.BCELoss, optim.Adam, DataLoader.
A Value engine in one file, layers in another, backprop in a third, optimizers somewhere else. To train a network today you copy-paste from five lessons and wire the pieces by hand. That is exactly the problem frameworks exist to solve — and they are not magic. They are organizational patterns you can build yourself.
Count what Phase 3 has given you: the perceptron (Lesson 01), multi-layer networks (02), backpropagation (03), activations (04), losses (05), optimizers (06), regularization (07), initialization (08), schedules (09). Ten lessons of working code — and no shared vocabulary. Every experiment starts with copy, paste, rename, re-wire. The source puts it plainly: a Value class here, a training loop there, weight initialization in another file.
The fix is not more math. It is an interface. Agree that every trainable piece in the system implements the same three verbs and one flag, and suddenly the pieces compose: any layer can sit next to any other layer, any container can hold any layer, and one training loop can train any network.
class Module:
def __init__(self):
self.training = True
def forward(self, x):
raise NotImplementedError
def backward(self, grad):
raise NotImplementedError
def parameters(self):
return []
def train(self):
self.training = True
def eval(self):
self.training = FalseA framework needs five groups of components, and you already built every one of them. The map below is the source’s architecture diagram, made clickable: select a group to see what it promises, which classes implement it, and which lesson it came from.
Five component groups, four arrows, one loop. Click a box — or a name in the row below — to see what it promises and which lesson built it.
The source draws this same picture as a Mermaid diagram. The arrow worth following is the loop: the optimizer’s output is the model’s next input — everything else is preparation for that one circular path.
Make the copy-paste cost concrete. The source’s problem statement describes training a network by wiring five lessons together. Suppose each lesson’s code is about 100 lines and you need the useful half from five of them, plus ~60 lines of glue — the loop, the batch iteration, the prints:
That last line is the whole argument. The source builds the framework in ~500 lines precisely so that each experiment afterwards is a dozen lines of intent — architecture, loss, optimizer, loop — with no duplicated plumbing to keep in sync.
Everything built in the next six chapters is a port of the source’s main.py: same classes, same method names, same order of operations. Where the browser labs use a simplified teaching model — no BatchNorm in the TypeScript port, a seeded PRNG instead of Python’s Mersenne Twister — the prose says so.
The Module contract is the reason a dropout layer, a linear layer and an entire eight-module network can be used interchangeably. Five methods, no exceptions. Everything else in this chapter is about what each layer puts behind that interface — and the single trick that makes gradients possible: caching.
forward(x) computes the output. For a Linear layer that is z = W·x + b: every output neuron is a weighted sum of the inputs plus its bias. backward(grad) receives the gradient of the loss with respect to this layer’s output and returns the gradient with respect to its input, while accumulating the gradients for its own weights and biases. parameters() returns those scalars in a flat list the optimizer can walk. train() and eval() flip a flag that stochastic layers read.
That is the whole vocabulary. The interesting engineering is inside forward, where every layer stores exactly what its backward will need — because a backward pass cannot re-run a forward pass; it must reuse the values that were seen.
Take the smallest useful layer, Linear(2, 1): two inputs, one output. Weigh it with w₁ = 0.5, w₂ = −0.25, bias b = 0.1, and feed it x = [2, −1]. In code that is the loop below; by hand it is three multiplications:
The weight initialization is also fixed by the contract. The source draws each weight from a Gaussian with standard deviation sqrt(2 / fan_in) — He initialization from Lesson 08: fan_in = 2 gives std 1.0 and weights like the ones above; fan_in = 16 gives 0.354, and fan_in = 784 gives 0.0505, because a wide input layer must not blow up the pre-activations with 784-term sums.
The activation modules are where caching becomes unavoidable. A ReLU cannot know a unit was on unless the forward pass wrote it down: its backward is grad × mask with mask = 1 where z > 0 else 0. Sigmoid saves its own output because its derivative is p·(1 − p) — a function of the output, not the input. Tanh saves its output for 1 − p². And Dropout saves a random mask whose non-zero entries are scaled by 1/(1 − p):
BatchNorm is the module that makes the training flag structural rather than cosmetic. In training mode it normalizes each feature across the batch — subtract the batch mean, divide by the batch standard deviation, then scale and shift with learned γ and β — and folds the batch statistics into running averages. In eval mode it stops looking at the batch entirely and uses those running averages. Note the split: γ and β are parameters (the optimizer updates them), while the running averages are buffers — state that must survive into inference and must never be touched by the optimizer.
class Linear(Module):
def __init__(self, fan_in, fan_out):
super().__init__()
std = math.sqrt(2.0 / fan_in) # He init, Lesson 08
self.weights = [[random.gauss(0, std) for _ in range(fan_in)]
for _ in range(fan_out)]
self.biases = [0.0] * fan_out
self.weight_grads = [[0.0] * fan_in for _ in range(fan_out)]
self.bias_grads = [0.0] * fan_out
self.input = None # the cache
def forward(self, x):
self.input = x # needed by backward: dW = grad·xᵀ
output = []
for i in range(self.fan_out):
val = self.biases[i]
for j in range(self.fan_in):
val += self.weights[i][j] * x[j]
output.append(val)
return output
def backward(self, grad):
input_grad = [0.0] * self.fan_in
for i in range(self.fan_out):
self.bias_grads[i] += grad[i] # += accumulates across a batch
for j in range(self.fan_in):
self.weight_grads[i][j] += grad[i] * self.input[j]
input_grad[j] += grad[i] * self.weights[i][j]
return input_gradTen frames: five forward stages, then the same five modules replayed in reverse. Watch the caches appear during the forward pass — the backward pass reads nothing else. Switch to eval mode and the dropout mask disappears.
The readout is the contract, not an illustration: each line comes from calling the module’s real forward() or backward() in the framework port. The dropout mask uses 0 and 1/(1−p) = 2, which keeps the expected value of every activation equal to its training-time average.
Why must a Linear layer store its input x during the forward pass? What would its weight gradient look like if it didn't?
Sequential is the composite pattern in one class: it holds a list of modules, forwards data through them left to right, backwards through them right to left, concatenates their parameter lists, and cascades the training flag. Because it satisfies the same contract, a stack can be nested inside another stack and nothing downstream notices.
The source’s circle classifier is eight modules and four linear layers, written as one call. The forward pass is a relay: x enters, each module’s output becomes the next module’s input, and every module writes down its cache as it goes. The backward pass is the same relay run right to left — the gradient the loss hands over becomes each module’s grad in turn.
class Sequential(Module):
def __init__(self, *modules):
super().__init__()
self.modules = list(modules)
def forward(self, x):
for module in self.modules:
x = module.forward(x)
return x
def backward(self, grad):
for module in reversed(self.modules):
grad = module.backward(grad)
return grad
def parameters(self):
params = []
for module in self.modules:
params.extend(module.parameters())
return params
def train(self):
self.training = True
for module in self.modules:
module.train()
def eval(self):
self.training = False
for module in self.modules:
module.eval()Trace one sample through the source’s model. Shapes are written as (features,); a batch of 16 samples would carry a (16, features) prefix through every step — the batch dimension is never touched by any layer.
A second, larger check — a digit classifier for 28×28 images, with a 784-wide input, one hidden layer of 128 and 10 output logits:
That last line is the hidden cost of the Module contract. Caching is what makes backward possible without re-running forward, and it is why a 7-billion-parameter model training on a GPU needs many times the weights’ memory for the gradients, optimizer moments and cached activations. Framework design is memory design.
Adding a layer is picking a type and a width. The board updates the forward shape chain and the parameter count — the same walkthrough Sequential does at runtime, drawn out where you can see it.
the lesson's 4-layer classifier: 2→16→16→8→1, 465 parameters, trained on the circle dataset
The parameter rule never changes: a Linear with fan_in inputs and fan_out outputs owns fan_in×fan_out weights plus fan_out biases. An activation owns nothing — it only caches.
A Linear layer has fan_in = 784 and fan_out = 128. How many trainable scalars does it own, and how does the framework find them?
Every module already knows its local derivative. The backward pass just delivers the loss gradient to the last module, and each module multiplies its way back through the chain — reading the caches from the forward pass and accumulating gradients for its own parameters.
There are only three things a module can do with an incoming gradient: use it (multiply by a local derivative), keep it (add the parameter gradient to the accumulator), and pass it (return the input gradient). Read the ReLU and Sigmoid code below with that sentence in mind — neither has parameters, so both only multiply and pass.
The subtle part is the multiplier. ReLU’s local derivative is 1 where the unit was on and 0 where it was off; Sigmoid’s is p·(1 − p); Linear’s is a matrix — the weights themselves for the input gradient, and the cached input for the weight gradients. The chain rule is everywhere, but each module only ever needs its own tiny piece.
class ReLU(Module):
def forward(self, x):
self.mask = [1.0 if v > 0 else 0.0 for v in x] # who fired
return [max(0.0, v) for v in x]
def backward(self, grad):
return [g * m for g, m in zip(grad, self.mask)] # dead units: 0 gradient
class Sigmoid(Module):
def forward(self, x):
self.output = []
for v in x:
v = max(-500, min(500, v)) # keep exp() finite
self.output.append(1.0 / (1.0 + math.exp(-v)))
return self.output
def backward(self, grad):
return [g * o * (1 - o) for g, o in zip(grad, self.output)] # g·p·(1−p)Use the tiny network from the composer: Linear(2,3) → ReLU → Linear(3,1) → Sigmoid, with one sample x = [1, 2] and target t = 1. The weights below are chosen so the first hidden unit is off — the dead-unit case every framework user eventually meets.
Check the multipliers for yourself: starting from dL/dz₂ = −0.4638, unit 1’s share is −0.4638 × W₂[1] = −0.4638 × (−0.4) = +0.1855 through the ReLU gate 1 — and that 0.1855 is exactly the bias gradient db₁[1]. The dead unit’s entire weight row receives 0.000 — not approximately zero, exactly zero, because one factor in every product is the mask. And note what the sigmoid + BCE pair collapsed to: dL/dz₂ = p − t = −0.4638. That tidy identity is why frameworks pair a sigmoid output with a binary cross-entropy loss.
Every forward operation is written down — the caches are the tape. Backward is the tape replayed in reverse, each recorded value multiplying its way into a gradient.
The first hidden unit was off (z1 = −0.4 → h = 0), so its ReLU mask is 0 and every gradient routed through it is exactly zero — the unit’s weights receive 0.000 this step. That is not a bug; it is the chain rule being honest.
A ReLU unit receives z = −0.4 in the forward pass. Its output is 0. What gradient flows back through it — and what do its weights receive this step?
A loss object turns predictions and targets into one scalar, then hands back the gradient that seeds the backward pass. An optimizer object holds a flat list of parameters — and no idea what architecture it is updating. Two small classes, and the loop has everything it needs.
The source’s loss classes use the same two-call shape as the modules: __call__ computes the loss and stores the prediction and target, then backward() returns the gradient with respect to the prediction. For MSE, that gradient is 2(p − t) per output, divided by the number of outputs so the loss is an average rather than a sum:
The optimizer is even smaller. It receives model.parameters() — a flat list of (value array, index, gradient array) triples, exactly the shape Linear returns — and applies one update rule to each entry. SGD multiplies the gradient by the learning rate; Adam keeps two running statistics per parameter and divides the cold start out. The optimizer never asks what a parameter is: it is blind by design, which is why swapping SGD for Adam is a one-line change.
class SGD:
def __init__(self, parameters, lr=0.01):
self.params = parameters # flat list of (container, i, j, grads)
self.lr = lr
def step(self):
for container, i, j, grad_container in self.params:
if j is not None: # a weight matrix entry
container[i][j] -= self.lr * grad_container[i][j]
else: # a bias
container[i] -= self.lr * grad_container[i]
def zero_grad(self):
for container, i, j, grad_container in self.params:
if j is not None:
grad_container[i][j] = 0.0
else:
grad_container[i] = 0.0The smallest network with a weight matrix: Linear(2, 1) with w₁ = 0.5, w₂ = −0.25, b = 0.1, fed x = [2, −1] and trained on the MSE loss against target 1. Everything is exact.
Both rules are exact for this input, and the lab below runs them for real: step as many times as you like, with an option that forgets zero_grad() so you can watch the accumulated gradients grow.
The smallest trainable model there is: z = w₁·2 + w₂·(−1) + b, MSE against target 1. Every number in the table comes from the lesson’s framework port — step it and check the arithmetic by hand.
| step | z | loss | grads (raw) | update | w₁ | w₂ | b |
|---|---|---|---|---|---|---|---|
| no steps yet — w₁ = 0.50, w₂ = -0.25, b = 0.10 | |||||||
raw grads are this step’s dL/dw; the update column is what the optimizer actually subtracted (gradients were cleared first).
With accumulate on, gradients add up instead of being cleared: the second step then moves by lr·(g₁ + g₂) instead of lr·g₂ — the arithmetic behind the zero_grad trap in chapter 04.
Everything so far is preparation for one loop. Batches arrive from the DataLoader, the forward pass caches, the loss names a scalar, the backward pass fills in gradients, the optimizer walks the flat parameter list. Repeat for a hundred epochs, then flip to eval mode and score.
The DataLoader solves two practical problems. A dataset may not fit in memory, so it hands you slices; and a fixed order lets a model exploit sequence rather than learn the pattern, so it reshuffles every epoch. The arithmetic is simple: the source draws 500 circle samples, splits 80/20, and trains on 400 of them. At batch size 16 that is ceil(400 / 16) = 25 batches per epoch — a number you can watch change in the lab below.
The loop order is the part worth memorizing: zero_grad → forward → loss → backward → step. Each call has one job and exactly one reason to be where it is. Clear first, or yesterday’s gradients join today’s step. Forward before loss, because the loss needs predictions. Backward before step, because the step reads gradients. And model.train() before the training loop, model.eval() before scoring — dropout and batch normalization behave differently in the two modes, and a test score taken in train mode is measuring the noise, not the model.
One honest wrinkle: inside each batch the source steps the optimizer per sample, not per batch. Batching organizes the epoch, but the updates are still per-sample stochastic gradient descent. It works — it is just noisier than true mini-batch updates, which is exactly what exercise 4 fixes. The lab computes the batch count as you move the slider, and shows (with a straight face) that the batch size does not change the trajectory until you make that change.
model = Sequential(
Linear(2, 16), ReLU(),
Linear(16, 16), ReLU(),
Linear(16, 8), ReLU(),
Linear(8, 1), Sigmoid(),
)
criterion = BCELoss()
optimizer = Adam(model.parameters(), lr=0.01)
loader = DataLoader(train_data, batch_size=16, shuffle=True)
model.train()
for epoch in range(100):
total_loss = 0
for batch_inputs, batch_targets in loader:
for x, t in zip(batch_inputs, batch_targets):
pred = model.forward(x) # forward: build the caches
loss = criterion(pred, t) # scalar loss
optimizer.zero_grad() # clear last step's gradients
grad = criterion.backward() # seed the backward pass
model.backward(grad) # replay the caches in reverse
optimizer.step() # w ← w − rule(gradient)
model.eval()
correct = 0
for x, t in test_data:
pred = model.forward(x) # no backward, no updates
if (pred[0] >= 0.5) == t[0]:
correct += 1
print(f"Test Accuracy: {correct / len(test_data) * 100:.1f}%")Both networks see the same data in the same order with the same optimizer objects. The dashed cyan line is the hand-written forward and backward pass; the solid accent line is Sequential + Modules + BCELoss. They lie exactly on top of each other — that is the point.
Honest surprise: changing the batch size does not move this curve at all. The source steps the optimizer once per sample, so batching only groups the epoch — the shuffled sample order is identical whatever the batch size. The batch-count readout changes, the trajectory does not. Exercise 4 turns this into true mini-batch updates, and then the batch size starts to matter.
You move optimizer.step() to before model.backward(grad). Nothing crashes. What actually happens over the next few steps?
You have a working framework. What separates it from PyTorch is not the abstraction — it is the answers underneath: who computes gradients, how parameters are found, where state lives, when gradients are cleared, and whether graph-building is eager or compiled. Read any new framework by asking those five questions.
1. Who computes gradients? In this lesson, you do — every module’s backward() is hand-written. In micrograd and PyTorch, a recorded graph does it: the forward pass writes down dependencies and backward() walks them. Either design gives the same numbers; autograd just removes the requirement that a human be right about every derivative.
2. How are parameters found? Your framework returns an explicit flat list from parameters(). PyTorch does the same thing but names each entry (named_parameters()), and that naming is load-bearing: it is how optimizers build parameter groups — the standard recipe skips weight decay for biases and normalization gains. Keras gathers layer.trainable_weights. The question to ask is always the same: who owns the list, and what happens if something is missing from it?
3. Where does state live? Three kinds: forward caches (what backward reads), persistent buffers (BatchNorm’s running mean and variance, which must survive into eval mode), and optimizer state (Adam’s m and v per parameter). Checkpointing only the weights loses the other two. When you add BatchNorm to your own framework, the running statistics go in the buffer category — parameters the optimizer must never update.
4. When are gradients cleared? Explicitly, by the training loop, which is what makes gradient accumulation possible. Keras hides it inside fit() and loses that flexibility; PyTorch keeps it visible, and this lesson kept it visible for the same reason — an optimizer that cleared its own gradients could never simulate a larger batch.
5. Eager or compiled? This framework is eager: each operation runs the moment it is called, which is what makes the debugger useful and the training loop readable. PyTorch is eager by default with an optional compiler; Keras/TensorFlow build a graph for speed; JAX traces a pure function of your step (Lesson 12). Eager loses some performance and wins all the debuggability.
Two honest differences worth naming. PyTorch operates on whole batches as tensors, so model(inputs) is one call where yours is a Python loop over samples; and PyTorch’s autograd means no one writes a backward() per layer, which is why its gradient code can be trusted at scale. Everything conceptual — composition, parameter registration, optimizer blindness, mode flags, zero_grad timing — transfers one-to-one.
The source’s punchline: PyTorch is not magic, it is this lesson with autograd bolted on. Switch tabs and watch the same 2→16→16→8→1 classifier change clothes.
model = Sequential(
Linear(2, 16), ReLU(),
Linear(16, 16), ReLU(),
Linear(16, 8), ReLU(),
Linear(8, 1), Sigmoid(),
)
criterion = BCELoss()
optimizer = Adam(model.parameters(), lr=0.01)
model.train()
for epoch in range(100):
for inputs, targets in loader:
for x, t in zip(inputs, targets):
optimizer.zero_grad() # 1. clear the accumulators
pred = model.forward(x) # 2. forward, caching along the way
loss = criterion(pred, t) # 3. scalar loss
grad = criterion.backward() # 4. gradient w.r.t. the prediction
model.backward(grad) # 5. each module multiplies by its local derivative
optimizer.step() # 6. update the flat parameter list| design question | micrograd | your framework | PyTorch | Keras |
|---|---|---|---|---|
| who computes gradients | autograd over scalars (the Value graph) | you do, in each module's backward() | autograd over tensors (the recorded graph) | autograd, hidden inside fit() |
| how are parameters found | walk the graph for nodes that need gradients | explicit flat list from parameters() | recursive model.parameters() — named_parameters() keeps names | layer.trainable_weights, gathered by the model |
| where does state live | Value.data and Value.grad per node | module caches + optimizer m and v arrays | tensor data and .grad + optimizer state dict | layer weights + optimizer slots |
| when are gradients cleared | not needed — each graph is fresh | zero_grad() before each step (yours to call) | optimizer.zero_grad(), usually set_to_none=True | once per batch, inside fit() |
| train versus eval | no distinction (no stochastic layers) | model.train() / model.eval() flag cascades | model.train() / model.eval(); with torch.no_grad() for scoring | the training argument on each call; fit() sets it |
When you read a new framework’s docs, ask this lab’s five questions. The answers tell you where the gradients come from, who owns them, and what will silently break when you forget to clear them.
Scientific honesty: the ~500-line framework in this lesson is a simplified teaching model, not a production library. The browser labs port its core (Linear, ReLU, Sigmoid, Dropout, Sequential, MSE, BCE, SGD, Adam, DataLoader) to TypeScript and omit BatchNorm, Tanh and the source’s tuple bookkeeping; the TypeScript PRNG is seeded mulberry32 rather than Python’s Mersenne Twister, so a browser run may differ from a Python run in the last decimal. The arithmetic that the lesson checks by hand is exact regardless of which port runs it.
The loop-order question and the zero_grad question are the two that separate “I built a framework” from “I can debug one.” Answer before looking; the explanations carry the numbers.
0 / 5 answered · 0 correct
01What three responsibilities does the Module abstraction have in a deep learning framework?
02Why does Sequential process modules in reverse order during the backward pass?
03Why is optimizer.zero_grad() a separate call instead of being done automatically inside step()?
04What is the correct order of operations in a training loop?
05What is the role of the DataLoader in the framework?
Click a card to swap the lazy description for what it actually means.
Four problems with exact numbers — add a combined softmax cross-entropy loss, put the learning rate on a schedule inside your optimizer, build checkpoint save/load, and replace per-sample stepping with true mini-batch accumulation. Try first; a worked answer is one click away.
Softmax: e² = 7.389, e¹ = 2.718, e⁰·¹ = 1.105; the sum is 11.212, so p = [0.6590, 0.2424, 0.0986]. Cross-entropy for target class 0: loss = −ln(0.6590) = 0.4170. The famous simplification is that softmax followed by cross-entropy has backward dL/dz = p − onehot(target) = [0.6590 − 1, 0.2424, 0.0986] = [−0.3410, 0.2424, 0.0986]. Two checks worth doing: the gradient entries sum to −0.3410 + 0.2424 + 0.0986 = 0 exactly (adding a constant to every logit cannot change the loss, so the gradients must cancel), and the largest gradient goes to the correct class, pushing its logit up while every wrong class is nudged down in proportion to its probability. In the class, that one line replaces the separate Sigmoid + BCE pair — and it is numerically safer, because the loss works with logits and never computes log(softmax) from a rounded probability.
Warmup is linear: lr(e) = 0.01·(e+1)/10 for e < 10, so lr(0) = 0.001, lr(4) = 0.005, lr(9) = 0.010. Cosine decay starts at epoch 10: lr(e) = 0.005·(1 + cos(π(e−10)/90)), giving lr(10) = 0.010, lr(55) = 0.005 (the cosine argument is π/2), and lr(100) = 0. The optimizer object needs one new method — SGD and Adam already hold lr as a field, so set_lr simply assigns it; the framework change is tiny because the schedule lives outside the model. What the model gets: tiny, cautious steps while Adam's moments are cold and the freshly initialized weights are least trustworthy (epochs 0–9); full-size steps through the middle of training when the gradients are informative; and a vanishing step at the end so the parameters settle instead of orbiting the minimum at radius ≈ lr. Compare against constant lr = 0.01 and the difference is usually a few points of final accuracy plus a smoother loss curve — exactly the experiment Lesson 09 describes, now running through your own loop.
The verification is exact and cheap: run a fixed batch of inputs through the original model and the loaded model in eval mode, and compare predictions. JSON round-trips IEEE-754 doubles exactly when you serialize with full precision (Python's json prints repr, JavaScript's JSON.stringify prints the shortest round-tripping form), so the check should report a maximum difference of exactly 0, not ≈ 0. If you see 1e-7-scale differences, you saved rounded decimals. The trap is that weights are not the whole model. A training checkpoint that only stores weights forgets: (1) the optimizer's state — Adam's m and v arrays, which are per-parameter and worth two extra floats per parameter; (2) BatchNorm's running_mean and running_var, which are buffers, not parameters, and are what eval mode actually uses; (3) the training/eval flag, because loading weights into a model still in train mode with dropout will give noisy, non-reproducible predictions. Any one of those omissions produces a model that looks like it loaded correctly and predicts differently — the classic 'my evaluation numbers changed after resume' bug.
Per-sample (what the source does): sample 1 gives p = 0.5, g = 2(p − t) = −1 for both w and b, so w ← 0.5 − 0.1·(−1)·1 = 0.6 and b ← 0.1. Sample 2 now runs at w = 0.6, b = 0.1: p = 0.6·2 + 0.1 = 1.3, g = 0.6, so g_w = 0.6·2 = 1.2, g_b = 0.6, and the second step lands at w = 0.6 − 0.12 = 0.48, b = 0.1 − 0.06 = 0.04. Accumulated: both gradients are computed at the same w = 0.5, b = 0 — g_w = −1 for sample 1 and g_w = 0 for sample 2 (p = 1.0 exactly, so the second sample is already correct) — the batch mean is g_w = −0.5, g_b = −0.5, and one step gives w = 0.5 + 0.05 = 0.55, b = 0.05. The two paths end at different points (0.48, 0.04) vs (0.55, 0.05), and neither is 'wrong': per-sample stepping takes many small, noisier updates; accumulation takes fewer, better-averaged ones. The trade is measured, not assumed — which is exactly why the source's exercise says to measure convergence rather than reason about it. One practical note: with accumulation, zero_grad() moves outside the sample loop and step() moves outside the batch loop; the gradient arrays fill up while you iterate.
You do not need to master these here. Each one gets a proper treatment in its own lesson; the one-line meaning is enough to keep reading. Orange dotted underlines in the prose point back to this list.
This lesson is a port of an open course. Everything here traces back to it — and the next step is running the code yourself.
Lesson text adapted from AI Engineering from Scratch (Phase 03, Lesson 10) and the Math Foundations Notebook reference build. The seven labs (framework map, Module-stack forward/backward stepper, layer composer, backward-tape animator, two-weight optimizer stepper, hand-rolled vs framework training dashboard, and the four-API comparison board), the TypeScript framework port, the 465-parameter and 101,770-parameter counts with their memory arithmetic, the hand-worked 2→3→1 forward and backward pass, the exact two-weight SGD and bias-corrected Adam steps, and the accumulation arithmetic in the exercises are original to this page. Every number shown is computed live by the labs or verified by hand in the prose; the source's own run is Python and its exact accuracies are not reproduced here.