Neural networks

A neural network is a chain of very simple steps: multiply the numbers by a grid of weights, add an offset, then bend the result so it is not a straight line, and repeat. This section is about why that structure can represent almost anything, and why the bend is the part that matters, since without it the whole stack collapses into a single multiplication no matter how many layers you give it.

Task                          Output layer size   Final activation   Loss
regression, one number        1                   none               mean squared error
regression, outliers present  1                   none               Huber or L1
binary classification         1                   none (logits)      binary cross-entropy with logits
multi-class, one label        number of classes   none (logits)      cross-entropy
multi-class, many labels      number of classes   none (logits)      binary cross-entropy per class
next token                    vocabulary size     none (logits)      cross-entropy

One layer#

A layer takes an input vector, multiplies it by a weight matrix, adds a bias, and applies a non-linear function. That is the entire operation, repeated.

It is worth doing once with numbers small enough to check by hand, because every shape error later is a failure to have this picture. Three input features, two output units. The weight matrix is therefore three rows by two columns, one column per output unit, and the bias has one entry per output unit.

import numpy as np

x = np.array([[1.0, 2.0, -1.0]])       # [1, 3]  one example, three features
W = np.array([[ 0.5, -0.2],            # [3, 2]  one column per output unit
              [ 1.0,  0.3],
              [-0.5,  0.8]])
b = np.array([0.1, -0.4])              # [2]     one per output unit

z = x @ W + b                          # [1, 3] @ [3, 2] -> [1, 2]
print(z)                               # [[ 3.1 -0.8]]
print(np.maximum(0, z))                # [[ 3.1  0. ]]

# unit 0:  1.0*0.5 + 2.0*1.0 + (-1.0)*(-0.5) + 0.1 =  3.1
# unit 1:  1.0*(-0.2) + 2.0*0.3 + (-1.0)*0.8 - 0.4 = -0.8  -> ReLU sets it to 0

One example was for legibility. In practice the first dimension is the batch, and nothing else changes: feeding 64 examples makes x shape [64, 3] and z shape [64, 2], with W and b untouched. That is why a batch costs almost nothing extra on a GPU. The same matrix multiply is issued once instead of 64 times.

def layer(x, W, b):
    return relu(x @ W + b)

# A three-layer network is that, three times
h1 = layer(x, W1, b1)
h2 = layer(h1, W2, b2)
out = h2 @ W3 + b3          # no activation on the output layer

input      [64,  784]     a batch of 64 flattened 28x28 images
W1         [784, 256]     -> h1 [64, 256]     784*256 = 200,704 weights
W2         [256, 128]     -> h2 [64, 128]     256*128 =  32,768 weights
W3         [128,  10]     -> out [64, 10]     128*10  =   1,280 weights
                                              plus 394 biases

logits     [64, 10]       one score per class per example

Why the non-linearity matters#

Without it, stacking layers is pointless. Two matrix multiplies in a row are equivalent to one matrix multiply, so a hundred linear layers can represent exactly what one linear layer can. The non-linearity between them is what makes depth mean anything.

x @ A @ B  =  x @ (A @ B)  =  x @ C

Three layers collapse to one. Add ReLU between them and they
no longer collapse, which is where the expressive power comes from.

That claim is usually asserted. It takes six lines to watch it happen, and the difference between having read it and having run it is what the follow-up question is for.

rng = np.random.default_rng(0)
A = rng.normal(0, 1, (4, 5))
B = rng.normal(0, 1, (5, 3))
C = rng.normal(0, 1, (3, 2))
X = rng.normal(0, 1, (100, 4))

deep   = ((X @ A) @ B) @ C           # three layers, no activation
single = X @ (A @ B @ C)             # one 4x2 matrix
print(np.abs(deep - single).max())   # 3.55e-15  -> the same, to floating point

relu = lambda t: np.maximum(0, t)
deep_relu = relu(relu(X @ A) @ B) @ C

# Fit the best possible single linear map to what the ReLU network computes.
M, *_ = np.linalg.lstsq(X, deep_relu, rcond=None)
print(np.abs(deep_relu - X @ M).max())   # 1.8619  -> no linear map comes close

The second half is the part that makes the point. It does not merely fail to find the equivalent matrix; it fits the best one that exists by least squares, and it is still wrong by 1.86 on outputs of roughly that size. There is no single matrix to find. Adding ReLU moved the network into a class of functions that no amount of linear algebra can reach.

Activations#

The choice is nearly always made for you by where in the network you are, and the list below is in the order you will meet them rather than by merit.

Activation   Formula              Range        Use it
ReLU         max(0, x)            0 to inf     hidden layers, the default, anything convolutional
GELU         x * P(Z <= x)        -0.17 to inf hidden layers of transformers
SiLU / Swish x * sigmoid(x)       -0.28 to inf hidden layers, common in recent vision and LLMs
Leaky ReLU   max(0.01x, x)        -inf to inf  hidden layers, when ReLUs are dying
Tanh         (e^2x - 1)/(e^2x + 1) -1 to 1     rarely now; recurrent gates and some output heads
Sigmoid      1 / (1 + e^-x)       0 to 1       output only, binary and multi-label probabilities
Softmax      normalised exp       sums to 1    output only, one-label multi-class

Sigmoid and tanh saturate: past roughly plus or minus five the output stops moving, so its gradient goes to nearly zero and the weights feeding it stop being updated. Since backpropagation multiplies these gradients together on the way back, a stack of saturated sigmoids annihilates the signal before it reaches the early layers. This is the vanishing gradient problem, and it is the concrete reason deep networks were not trainable before ReLU.

def sigmoid(t): return 1 / (1 + np.exp(-t))

for v in (0.0, 2.0, 5.0, 10.0):
    s = sigmoid(v)
    print(v, round(s, 6), round(s * (1 - s), 8))

# x= 0.0   output 0.500000   gradient 0.25000000   the best it ever gets
# x= 2.0   output 0.880797   gradient 0.10499359
# x= 5.0   output 0.993307   gradient 0.00664806
# x=10.0   output 0.999955   gradient 0.00004540   effectively dead

print((sigmoid(5) * (1 - sigmoid(5))) ** 10)     # 1.69e-22 through ten such layers

ReLU has gradient exactly 1 wherever it is active, so it multiplies nothing away, and it costs a comparison rather than an exponential. Its own failure is the mirror image: a unit whose input is always negative outputs zero, has zero gradient, and never recovers. That is the dying ReLU problem, and a learning rate set too high is the usual cause, because one enormous update pushes the unit's bias far negative and it never comes back.

Leaky ReLU fixes it bluntly by keeping a small slope for negatives. GELU and SiLU fix it smoothly, and are what transformers use: both are close to ReLU for large inputs but curve gently through zero, so a unit sitting slightly negative still receives a small gradient and can return. That is most of the practical difference, and it is worth perhaps a fraction of a per cent. Nobody has lost a job over picking GELU instead of SiLU.

Loss functions#

The loss is the only place where you say what you mean by a good answer. Everything else in training is machinery for lowering whatever number the loss produces, so choosing it wrongly means the machinery works perfectly towards the wrong target.

Mean squared error is for regression, and it squares, so a prediction ten units off counts a hundred times a prediction one unit off. That is right when large errors are genuinely much worse, and wrong when the data contains outliers, where it drags the whole model towards a handful of bad rows. Huber loss is the standard compromise: squared for small errors, linear for large ones, so outliers are noticed but not obeyed.

Cross-entropy is for classification, and it is the negative log of the probability assigned to the correct answer. Being right with probability 1 costs zero. Being wrong costs more and more without limit as the confidence rises, which is the behaviour you want and is precisely what squared error fails to provide.

for p in (0.9, 0.5, 0.1, 0.01, 0.001):
    print(p, round((1 - p) ** 2, 5), round(-np.log(p), 4))

# predicted probability of the TRUE class
#   p       MSE        cross-entropy
# 0.900   0.01000        0.1054
# 0.500   0.25000        0.6931
# 0.100   0.81000        2.3026
# 0.010   0.98010        4.6052
# 0.001   0.99800        6.9078

MSE is bounded above by 1 here, so once a classifier is confidently wrong there is almost no gradient left to correct it. Cross-entropy is unbounded, so the more confidently wrong it is the harder it is pushed back. That is the whole argument for using it, and it is a better answer than reciting that cross-entropy is for classification.

Loss                              Task                             Notes
mean squared error                regression                       assumes Gaussian noise, hates outliers
Huber / smooth L1                 regression with outliers         one threshold parameter to set
cross-entropy                     one label out of many            feed it logits, never softmax output
binary cross-entropy with logits  yes-or-no, or multi-label        numerically stable version of the above
contrastive / InfoNCE             embeddings                       learns that similar things sit close
CTC                               sequences without alignment      speech and handwriting

Deep networks did not become trainable because somebody stacked more layers. Three unglamorous components made it possible: initialising weights so the signal neither explodes nor vanishes on the way through, normalising activations so each layer sees inputs in a stable range, and residual connections that give gradients a direct path back. This section covers what each does and why. Get initialisation wrong and a deep network simply does not train, and the failure hands you nothing to work with.

Initialisation#

Weights initialised too large make activations explode through the layers; too small and they vanish. Schemes such as Xavier and He scale the initial weights by the layer's size so the signal keeps a stable magnitude at the start of training. Get this wrong and a deep network does not train at all.

The reasoning is one line of arithmetic. A unit sums n inputs, each multiplied by a weight, so its variance is roughly n times the variance of one weight. Do that fifty times and the signal is multiplied by that factor fifty times over, which either explodes or disappears unless the factor is close to 1. So the weights are drawn with variance about 1 over n, which is Xavier. ReLU then throws away the negative half, halving the variance again, so the correction is 2 over n, which is He.

import numpy as np

def propagate(name, std, depth=50, width=256):
    rng = np.random.default_rng(0)
    h = rng.normal(0, 1, (128, width))
    for i in range(depth):
        W = rng.normal(0, std(width), (width, width))
        h = np.maximum(0, h @ W)               # ReLU network, 50 layers deep
        if i in (0, 9, 49):
            print(name, i + 1, f"{np.sqrt((h ** 2).mean()):.3e}")

propagate("too small (0.01)",  lambda w: 0.01)
propagate("too large (0.1)",   lambda w: 0.1)
propagate("Xavier sqrt(1/n)",  lambda w: (1.0 / w) ** 0.5)
propagate("He     sqrt(2/n)",  lambda w: (2.0 / w) ** 0.5)

#                       layer 1     layer 10    layer 50
# too small (0.01)    1.125e-01   3.555e-10   2.388e-48   nothing reaches the output
# too large (0.1)     1.125e+00   3.555e+00   2.388e+02   and keeps growing
# Xavier sqrt(1/n)    7.033e-01   3.233e-02   1.486e-08   decays: the ReLU half was not accounted for
# He     sqrt(2/n)    9.947e-01   1.035e+00   4.985e-01   stable through all fifty

Read the failure rows carefully, because they describe what you will actually see. Too small gives a forward signal of 1e-48, so every output is effectively zero, the loss sits exactly at its initial value, and nothing moves; it looks like a bug in the data loader. Too large gives activations in the hundreds, then thousands, and the first NaN arrives within a few steps. Neither hands you an error message that mentions initialisation.

Xavier is the instructive row. It is not wrong, it is correct for a symmetric activation such as tanh, and used with ReLU it loses a factor of the square root of 2 per layer, which over fifty layers is the difference between 0.5 and 1e-8. This is why frameworks want to know your activation: PyTorch's kaiming_normal_ takes a nonlinearity argument, and passing the wrong one is a silent, slow failure rather than a loud one.

One case is not about scale at all. If every weight starts at the same value, including zero, then every hidden unit in a layer computes the same thing, receives the same gradient, and is updated identically forever. A thousand-unit layer behaves as one unit, permanently. Randomness is what breaks that symmetry, and it is a separate requirement from getting the magnitude right.

rng = np.random.default_rng(0)
X = rng.normal(0, 1, (32, 4)); y = rng.normal(0, 1, 32)

def first_gradient(W1, W2):
    h = np.maximum(0, X @ W1)
    d_out = (h @ W2 - y[:, None]) / len(y)
    return X.T @ ((d_out @ W2.T) * (h > 0))

same = first_gradient(np.full((4, 6), 0.5), np.full((6, 1), 0.5))
rand = first_gradient(rng.normal(0, 0.5, (4, 6)), rng.normal(0, 0.5, (6, 1)))
print(np.round(same[0], 4))   # [0.0819 0.0819 0.0819 0.0819 0.0819 0.0819]  identical
print(np.round(rand[0], 4))   # [-0.0284 0.0124 0.0156 -0.0488 0.0163 -0.0234]

Normalisation#

Initialisation only fixes the first step. Once the weights start moving, the distribution of activations drifts, and every layer is chasing a moving target set by the layers below it. Normalisation pins that down during training: subtract the mean, divide by the standard deviation, then apply a learned scale and shift so the layer can still choose any range it wants.

def normalise(h, axis, gamma, beta, eps=1e-5):
    mean = h.mean(axis=axis, keepdims=True)
    var = h.var(axis=axis, keepdims=True)
    return gamma * (h - mean) / np.sqrt(var + eps) + beta

h = np.random.default_rng(0).normal(3.0, 7.0, (32, 128))    # [batch, features]

batch_normed = normalise(h, axis=0, gamma=1.0, beta=0.0)    # down the batch, per feature
layer_normed = normalise(h, axis=1, gamma=1.0, beta=0.0)    # across features, per example

print(round(batch_normed.mean(0).max(), 6), round(batch_normed.std(0).min(), 4))   # ~0.0  ~1.0
print(round(layer_normed.mean(1).max(), 6), round(layer_normed.std(1).min(), 4))   # ~0.0  ~1.0

• Batch normalisation rescales activations using statistics across the batch. Effective in vision, and awkward with small batches or sequences.
• Layer normalisation rescales across the features of each example independently. The standard in transformers, since it does not depend on batch composition.

Normalising down the batch means a prediction depends on which other examples happened to be in the batch with it. During training that is a mild regulariser. At inference there is no batch, so batch norm keeps a running average of the statistics it saw and uses those instead, which is a second set of numbers that must be saved with the model and is why model.eval() changes its behaviour. With a batch of size 2 those statistics are noise, and with variable-length sequences the padding contaminates them. Layer norm has neither problem because it never looks outside a single example, which is the entire reason transformers use it.

RMSNorm is the current default in large language models and is layer norm with the mean subtraction removed, keeping only the division by the root mean square. It works about as well and is measurably cheaper, which matters when it runs twice per block across eighty blocks.

Residual connections#

Add a layer's input to its output, so the layer learns a change rather than a whole new representation. This gives gradients a direct path back through the network, which is what made very deep networks trainable. Almost every modern architecture uses them.

The argument is about what the chain rule does to a product. Without a residual, the gradient reaching layer 1 is the product of forty Jacobians, and a product of forty numbers that are each a little below 1 is nothing at all. With a residual, each term becomes 1 plus something, so there is always a path back that multiplies by 1. Forty layers deep, that is the difference between a first layer that trains and one that does not.

import torch, torch.nn as nn

class Net(nn.Module):
    def __init__(self, depth=40, width=128, residual=False):
        super().__init__()
        self.layers = nn.ModuleList([nn.Linear(width, width) for _ in range(depth)])
        self.residual = residual

    def forward(self, x):
        for lin in self.layers:
            out = torch.relu(lin(x))
            x = x + out if self.residual else out       # the entire change
        return x

for residual in (False, True):
    torch.manual_seed(0)
    net = Net(residual=residual)
    net(torch.randn(64, 128)).sum().backward()
    g = [l.weight.grad.norm().item() for l in net.layers]
    print(residual, f"{g[0]:.3e}", f"{g[39]:.3e}", f"{g[39] / g[0]:.3e}")

# residual=False   layer 0: 3.891e-13   layer 39: 2.072e+02   ratio 5.3e+14
# residual=True    layer 0: 6.565e+06   layer 39: 2.796e+07   ratio 4.3e+00

A ratio of 5e14 means the first layer is not training slowly; it is not training. Any learning rate large enough to move it would have destroyed layer 39 long before. With the skip connection the ratio is about 4, so one learning rate suits the whole network, and that is what made depth a parameter you can increase rather than a wall.

class Block(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.ff = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim))

    def forward(self, x):
        return x + self.ff(self.norm(x))     # residual: input plus change

This is where the pieces meet. You write a two-layer network with nothing but NumPy: the forward pass that turns inputs into predictions, the backward pass that works out how much each weight contributed to the error, and the loop that repeats both until the thing learns. Nobody builds production models this way. The point is that after doing it once, a PyTorch model stops being magic and becomes a stack of things you have already written yourself.

Three spirals#

import numpy as np
rng = np.random.default_rng(0)

N, K = 100, 3                                    # points per class, classes
X = np.zeros((N * K, 2)); y = np.zeros(N * K, dtype=int)
for k in range(K):
    r = np.linspace(0.0, 1.0, N)
    angle = np.linspace(k * 4, (k + 1) * 4, N) + rng.normal(0, 0.2, N)
    X[k * N:(k + 1) * N] = np.c_[r * np.sin(angle), r * np.cos(angle)]
    y[k * N:(k + 1) * N] = k

That claim is checkable in two lines, and worth checking, because the whole exercise is pointless on data a straight line already handles. Fitting scikit-learn's logistic regression to these 300 points scores 0.533. The network below reaches 0.993 on the same data with no feature engineering at all; the hidden layer is doing the work that a person would otherwise do by inventing polar coordinates.

Forward#

H = 100
W1 = rng.normal(0, 0.01, (2, H)); b1 = np.zeros(H)
W2 = rng.normal(0, 0.01, (H, K)); b2 = np.zeros(K)

def forward(X):
    h = np.maximum(0, X @ W1 + b1)                      # [n, H]  hidden layer, ReLU
    logits = h @ W2 + b2                                # [n, K]
    e = np.exp(logits - logits.max(axis=1, keepdims=True))
    p = e / e.sum(axis=1, keepdims=True)                # softmax
    return h, p

def loss_of(p, y):
    return -np.log(p[np.arange(len(y)), y]).mean()      # cross-entropy

Two details in there are not decoration. Subtracting the row maximum before the exponential changes nothing mathematically, because the constant cancels in the division, but without it a logit of 800 becomes inf and the softmax returns NaN; every framework does this internally and it is why you are told to hand the loss raw logits. And the indexing in loss_of picks out the probability of the correct class for each row in one operation, which is what cross-entropy is: the negative log of the probability you assigned to the right answer.

Before training anything, check the loss at step 0. The weights are near zero, so all three classes get roughly equal probability, so the loss should be the negative log of one third, which is 1.0986. It prints 1.0986. If it had printed 7, the labels or the initialisation would be wrong, and that is a thirty-second check rather than an hour-long one.

Backward#

def backward(X, y, h, p):
    n = len(y)
    d_logits = p.copy()
    d_logits[np.arange(n), y] -= 1                      # p - y, from the calculus section
    d_logits /= n
    dW2 = h.T @ d_logits;  db2 = d_logits.sum(axis=0)
    dh = d_logits @ W2.T                                # send the error back through W2
    dh[h <= 0] = 0                                      # ReLU passes gradient only where it was active
    dW1 = X.T @ dh;        db1 = dh.sum(axis=0)
    return dW1, db1, dW2, db2

Each line is the chain rule applied to one operation, in reverse order. The shape of each gradient matches the shape of the thing it is the gradient of, which is the quickest check that a line is right.

Three of those lines carry the whole idea and are worth saying out loud. The first, that the gradient of cross-entropy with respect to the logits is simply the predicted probabilities minus the one-hot target, is the single most useful fact in this section: softmax and cross-entropy are derived together precisely so that everything cancels into a subtraction. It is also why frameworks fuse them into one operation and refuse to take a softmax you applied yourself.

The second is dh, which sends the error backwards through W2 by multiplying by its transpose. Forward, the weights map hidden units to outputs; backward, the same numbers map output errors to hidden units, and the transpose is nothing more cunning than that. The third is the ReLU line: a unit that was inactive contributed nothing to the output, so it receives no blame. That one line is where the next section's bug lives.

Check the gradients#

A hand-written backward pass is a derivation, and derivations have errors. The check does not need any calculus: nudge one parameter up by a tiny amount, nudge it down, see how much the loss actually moved, and compare that with what your formula claimed. If they disagree, the formula is wrong.

def numerical(param, eps=1e-5):
    g = np.zeros_like(param)
    for i in np.ndindex(param.shape):
        old = param[i]
        param[i] = old + eps; plus = loss_of(forward(X)[1], y)
        param[i] = old - eps; minus = loss_of(forward(X)[1], y)
        param[i] = old
        g[i] = (plus - minus) / (2 * eps)
    return g

h, p = forward(X)
analytic = backward(X, y, h, p)[2]                      # dW2
num = numerical(W2)
print(np.abs(analytic - num).max() / (np.abs(analytic).max() + 1e-12))   # 9.03e-09

Relative error rather than absolute, because a gradient of 1e-7 and a gradient of 1e7 are both legitimate and an absolute threshold would pass one and fail the other. Below about 1e-7 the derivation is right. Above 1e-2 it is wrong. In between, suspect the kink described below before suspecting yourself.

Now delete the ReLU masking line and run the same check against dW1.

dW1, ReLU mask present    relative error 5.02e-03
dW1, ReLU mask removed    relative error 5.10e-01     <- a hundred times worse

Then train both for 5,000 steps anyway:
  with the mask       final loss 0.028   accuracy 0.993
  without the mask    final loss 0.572   accuracy 0.793     <- it still learns

That is why this check exists. The broken network does not crash, does not produce NaN, and does not sit at chance; it trains to 79 per cent, which on an unfamiliar dataset is entirely plausible and would be shipped. Nothing but the gradient check distinguishes it from the correct one.

Train#

No optimiser, no momentum, no schedule: subtract the learning rate times the gradient from every parameter, five thousand times. A learning rate of 1.0 would be wild for a deep network and is fine here because the network is two layers on 300 points, and there is no mini-batching either, since the whole dataset fits in one matrix.

lr = 1.0
for step in range(5001):
    h, p = forward(X)
    grads = backward(X, y, h, p)
    for param, g in zip((W1, b1, W2, b2), grads):
        param -= lr * g
    if step % 1000 == 0:
        print(step, round(loss_of(p, y), 3), round((p.argmax(1) == y).mean(), 3))

# step     loss    accuracy
#    0    1.099      0.463      log(3): the model knows nothing
# 1000    0.134      0.947
# 2000    0.070      0.983
# 3000    0.047      0.987
# 4000    0.035      0.993
# 5000    0.028      0.993      logistic regression on the same data: 0.533

The curve is the ordinary shape: most of the progress in the first thousand steps, then a long tail buying decreasing amounts. Accuracy stops at 0.993 while the loss keeps falling, which is the model becoming more confident about points it already had right rather than fixing the two it has wrong. Those two sit where the spiral arms meet and are genuinely ambiguous.

A training run emits a stream of numbers, and most of what goes wrong in machine learning shows up there long before it shows up in a final score. This section is about reading that stream: what a healthy loss curve looks like, what a learning rate set too high does to it, how overfitting announces itself, and what gradient statistics tell you when the curve alone is ambiguous. It is the difference between waiting an hour for a run to fail and killing it after two minutes knowing what to change.

What healthy looks like#

Training loss falls quickly at first and then more slowly. Validation loss follows it closely, then flattens a little above it. Both are smooth apart from mini-batch noise.

Every curve below came out of the same twenty-line script, with one number changed each time. It is worth producing them yourself once, because a description of a diverging loss is much less memorable than having watched one.

import torch, torch.nn as nn, torch.nn.functional as F

X = torch.randn(2000, 20); w = torch.randn(20, 3)
y = (X @ w).argmax(1)                                  # a learnable problem, by construction
Xtr, ytr, Xva, yva = X[:1500], y[:1500], X[1500:], y[1500:]

def run(lr, steps=400, n=1500):
    torch.manual_seed(0)
    model = nn.Sequential(nn.Linear(20, 64), nn.ReLU(), nn.Linear(64, 3))
    opt = torch.optim.SGD(model.parameters(), lr=lr)
    for s in range(steps):
        loss = F.cross_entropy(model(Xtr[:n]), ytr[:n])
        opt.zero_grad(); loss.backward()
        gnorm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1e9)   # measure, do not clip
        opt.step()
        if s in (0, 50, 100, 200, 399):
            with torch.no_grad():
                val = F.cross_entropy(model(Xva), yva).item()
                dead = (model[0](Xtr[:n]).max(0).values <= 0).float().mean().item()
            print(s, round(loss.item(), 4), round(val, 4), round(gnorm.item(), 3), round(dead, 3))

run(lr=0.5)      HEALTHY
step    train      val    grad norm   dead units
   0   1.1429   1.0489      0.485        0%        starts at log(3) = 1.0986
  50   0.1412   0.1646      0.071        0%
 100   0.0770   0.1052      0.036        0%
 399   0.0204   0.0641      0.011        0%        val a little above train, both falling

run(lr=0.001)    LEARNING RATE TOO LOW
   0   1.1429   1.1485      0.485        0%
 200   1.1011   1.1065      0.432        0%
 399   1.0672   1.0728      0.397        0%        400 steps to move 0.08. A straight line.

run(lr=50)       LEARNING RATE TOO HIGH
   0   1.1429  30.9919      0.485        0%
  50   318239   2.07e+17    59.157     65.6%       one huge step, then wreckage
 100   5.5275   2.07e+17     0.683     65.6%       train recovers, val never does
 399  13.9810   2.07e+17     0.644     65.6%       two thirds of the units are dead

The third run is the one worth studying, because its training loss comes back down to single digits and would look survivable in a plot that started at step 100. It is not survivable. One step at lr=50 pushed 65.6 per cent of the hidden units to a permanently negative input, and a ReLU there outputs zero with zero gradient forever. The network is now a third of the size you built, the damage happened in a single step, and nothing in the logs says so unless you measure dead units.

Overfitting looks different again, and cannot be seen at all without a validation split. The same script on 40 training rows instead of 1,500:

run(lr=0.5, steps=2000, n=40)     OVERFITTING
step     train      val
   0    1.1979   1.0556
 100    0.0065   1.2157      train already essentially zero; val has turned
 400    0.0011   1.4405
1000    0.0004   1.5850
1999    0.0002   1.6929      still improving on train, still worsening on val

Note that the training loss never misbehaves. It falls smoothly to 0.0002 and would be reported as an excellent run by anyone logging one number. The entire signal is in the second column, and the model worth shipping existed for about nine steps. This is the argument for checkpointing on best validation rather than at the end: the end of training is not where the best model is, and on a small dataset it is nowhere near.

A field guide#

Symptom                                            Likely cause                 Try
Loss jumps around, sometimes explodes              learning rate too high       lower it, add warmup, clip gradients
Loss falls painfully slowly in a straight line     learning rate too low        raise it by 3 to 10x
Loss flat from step 0                              bug, frozen parameters       overfit one batch, check gradients
Train falls, validation rises after a point        overfitting                  early stop, regularise, more data
Loss plateaus, then suddenly drops                 normal for some tasks        be patient before changing things
Loss becomes NaN                                   overflow, log(0)             lower LR, bfloat16, stable loss functions

Inside the network#

When the loss curve is ambiguous, the next numbers to look at are inside the model. Three are enough to separate most causes, and all three cost a line each.

• Dead ReLUs: units whose input is always negative output zero and never learn again. Track the share of zero activations per layer; a large learning rate is the usual cause.
• Vanishing gradients: gradient norms in early layers far smaller than in late layers. Residual connections, normalisation and good initialisation are the fixes.
• Exploding gradients: sudden spikes in gradient norm. Clipping caps them.

@torch.no_grad()
def probe(model, x):
    """Dead units per layer, and gradient norm per layer. Run it every few hundred steps."""
    h = x
    for i, layer in enumerate(model):
        h = layer(h)
        if isinstance(layer, nn.ReLU):
            print("layer", i, "dead fraction", round((h.max(0).values <= 0).float().mean().item(), 3))
    for name, p in model.named_parameters():
        if p.grad is not None:
            print(name, "grad norm", round(p.grad.norm().item(), 6))

What you measure                              What it means                  What to do
dead fraction rising over training            learning rate too high         lower it; the dead units do not come back
dead fraction high from step 0                initialisation, or a bad bias  check the init scheme and the activation
early-layer grad norm << late-layer           vanishing gradients            residuals, normalisation, He init
grad norm spikes by 100x on some batches      a few pathological examples    clip, and go and look at those batches
grad norm exactly 0.0 somewhere               detached graph, or frozen      check requires_grad and no_grad scopes
grad norm steady, loss flat                   the loss is wrong, not the model

Log these from the start#

Adding logging after a run goes wrong means running it again to find out why. All of this is cheap except the last, and the last is the one that catches the problems the numbers cannot.

• Training and validation loss.
• The learning rate at each step.
• The total gradient norm, and per-layer norms if something is odd.
• A few sample predictions at intervals, read by a person.
• Wall-clock time per step, which is how you notice the data loader has become the bottleneck.
• The step number of the best validation loss so far, and a checkpoint at it.