Training a Network

This is the whole of learning, in four moves: run the batch forward, score it, propagate the gradient back, then step every parameter against its gradient.

logits, z1, a1 = two_layer_forward(...)
P = softmax_rows(logits), loss = cross_entropy_from_logits(logits, labels)
output_delta = softmax_ce_backward(P, labels), then the four gradients
$\theta\leftarrow\theta-\eta\,\bar\theta$ for each of $W_1,b_1,W_2,b_2$

All the helpers are provided; the loop body is yours. Report the loss before the update.

Step 2 draws on two different forms of the same scores: the loss reads the raw logits, because a probability can underflow to 0.0 and $\log 0$ is $-\infty$ , while the backward pass in step 3 needs the probabilities P themselves.

Repeating the step on one fixed batch at $\eta=0.5$ :

after `steps` calls	reported `loss`
`1`	`0.7553`
`2`	`0.7326`
`3`	`0.7148`
`6`	`0.6744`
`26`	`0.5174`

The order is forced. Nothing can be updated until every gradient is computed, because $\bar W_1$ is derived from the current $W_2$ — update $W_2$ first and the first layer gets gradients for a network that no longer exists. So: all four gradients, then all four updates.

Reporting the pre-update loss is the standard convention, and it is why the sequence above starts at the loss of the untouched parameters. It also makes the number free: it was already computed on the way to the gradient.

Return new arrays rather than modifying the caller's in place — W1 - lr * dW1, not W1 -= lr * dW1. The caller may still be holding those parameters (to compare against, or to roll back), and a step that silently mutates them is a bug that surfaces far from its cause.

Python

def train_step(X: np.ndarray, labels: np.ndarray, W1: np.ndarray, b1: np.ndarray,
               W2: np.ndarray, b2: np.ndarray, lr: float):
    """One mini-batch training step: forward, loss, backward, update.

    X:      [B, D] input batch
    labels: [B]    integer class ids
    W1: [D, H], b1: [H], W2: [H, C], b2: [C]
    lr:     learning rate

    Steps, using the provided helpers:
      1. logits, z1, a1 = two_layer_forward(X, W1, b1, W2, b2)
      2. P = softmax_rows(logits); loss = cross_entropy_from_logits(logits, labels)
      3. output_delta = softmax_ce_backward(P, labels)
      4. dW1, db1, dW2, db2 = two_layer_backward(X, W1, z1, a1, W2, output_delta)
      5. subtract lr * gradient from each of the four parameters

    Returns (loss, W1, b1, W2, b2) where loss is measured BEFORE the update and
    the four parameters are the new ones. The caller's arrays are not modified.
    """
    logits, z1, a1 = two_layer_forward(X, W1, b1, W2, b2)
    P = softmax_rows(logits)
    loss = cross_entropy_from_logits(logits, labels)
    output_delta = softmax_ce_backward(P, labels)
    dW1, db1, dW2, db2 = two_layer_backward(X, W1, z1, a1, W2, output_delta)
    return loss, W1 - lr * dW1, b1 - lr * db1, W2 - lr * dW2, b2 - lr * db2

One mini-batch training step

Why weights can't all start the same