Deep networks

A deep network's gradient is a product with one factor per layer, and every fix in this lesson is a way of holding that factor near 1. This level builds the problem in a chain of ten numbers, watches it in a 30-layer network, then adds each fix and measures what it restores.

Picture a game of telephone along a line of 30 people. Each person repeats the message to the next, but everyone speaks at a quarter of the volume they heard. By the end of the line the message is silence. If instead everyone speaks 1.5× louder, the end of the line is a deafening roar. Training a deep network has exactly this problem, run backwards: the learning signal (the gradient, how much each weight should change; see primer.ml.neural_net) starts at the output and is passed back layer by layer, and each layer multiplies it by its own slope (derivative). Thirty multiplications by something below 1 is almost zero: the vanishing gradient. Thirty by something above 1 is enormous: the exploding gradient.

Worked example: a chain of ten one-number layers, each sitting at its steepest point.

flowchart RL
  L[Loss] -- "gradient 1" --> H10[layer 10]
  H10 -- "× 0.25" --> H9[layer 9]
  H9 -- "× 0.25" --> H8[layer 8]
  H8 -- "× 0.25 ... " --> H2[layer 2]
  H2 -- "× 0.25" --> H1["layer 1<br/>receives 0.25¹⁰ ≈ 1e-6"]

def gradient_at_input(w_l, slope, L=10):
    # ∂L/∂h_L: the loss hands the top layer 1
    grad = 1.0
    # Π over the layers: a running product
    for l in range(L):
        # × w_l φ'(z_l)
        grad *= w_l * slope
    return grad
# sigmoid at its steepest
f"{gradient_at_input(1, 0.25):.1e}"  # → '9.5e-07'
# linear, weight 1
gradient_at_input(1, 1)  # → 1.0
# linear, weight 1.5
round(gradient_at_input(1.5, 1), 1)  # → 57.7

chain_gradient builds that chain and backprops through it.

Why it matters this is why networks deeper than a handful of layers were considered untrainable for decades. Every fix below (better activations, careful initialization, residual connections, normalization) is a way of keeping the per-layer factor close to 1.

In a real layer, each neuron sums 64 inputs, so the multiplier per layer depends on three things together: the size of the weights, how many inputs each neuron adds up, and the activation's slope. Same telephone game, but now everyone in the line hears 64 people at once.

Worked example: a 30-layer network, 64 neurons per layer. The ratio of the gradient at the first layer to the gradient at the last:

import math
n_in, typical_slope = 64, 0.7
for sigma_w in (0.01, 1.0):
    # σ_w √n_in · typical φ'
    gain = sigma_w * math.sqrt(n_in) * typical_slope
    # per layer, then over 30 layers
    print(round(gain, 3), f"{gain ** 30:.0e}")  # → 0.056 3e-38 5.6 3e+22

In code: gradient_norms runs a 30-layer, 64-wide network forward and backward and returns the gradient size reaching every layer; first_to_last_gradient_ratio divides the first by the last to fill the table.

Why it matters you can't see this from the loss curve alone. A network whose early layers get no gradient still trains a little (the late layers learn), just badly. Plotting per-layer gradient norms is a standard diagnostic.

Think of a chain of 30 audio amplifiers. If each is set a little too quiet, the sound fades to nothing; a little too loud, and it distorts into noise. Set each so that what comes out is exactly as loud as what went in, and the music survives the whole chain. Initialization picks the random starting weights' size so that each layer passes on a signal of the same size, forward and backward.

Worked example:

He uses twice Xavier's variance because ReLU zeroes about half its inputs, throwing away half the signal's energy; the factor 2 puts it back.

flowchart LR
  A{Activation?} -->|ReLU / GELU| He["He: std = √(2 / fan_in)"]
  A -->|tanh / sigmoid / linear| X["Xavier: std = √(2 / (fan_in + fan_out))"]
  He & X --> S[Signal keeps its size<br/>layer after layer]

import math
n_in = 50
# Var(w) = 2 / n_in, for ReLU
var_w = 2 / n_in
# the variance, then the standard deviation
var_w, round(math.sqrt(var_w), 3)  # → (0.04, 0.2)
# Xavier for comparison: 100 in, 100 out
round(math.sqrt(2 / (100 + 100)), 3)  # → 0.1

In code: init_std returns the starting weight standard deviation for Xavier, He and two deliberately bad choices; forward_signal_rms measures the forward signal plotted above.

Why it matters every framework initializes this way by default (PyTorch's nn.Linear uses a Kaiming-style uniform init). Custom layers or deep stacks built without it can silently fail to train.

Picture a building where messages go up by stairs, one floor at a time, and at every landing someone might mumble. Add an express lift that runs the whole height, and the message always arrives intact; each floor adds its own notes to what the lift carries. A residual (or skip) connection is that express lift: each block computes a correction and adds it to its input, instead of replacing the input.

Worked example: ten blocks, each with slope 0.025 of its own.

flowchart TD
  IN[h] --> F["block F<br/>(layers, activation)"]
  F --> ADD((+))
  IN -- "skip: identity" --> ADD
  ADD --> OUT["h + F(h)"]

# each block's own slope
dF_dh = 0.025
plain, residual = 1.0, 1.0
for l in range(10):
    # h ← F(h): the slope is just ∂F/∂h
    plain *= dF_dh
    # h ← h + F(h): the slope is I + ∂F/∂h
    residual *= 1 + dF_dh
f"{plain:.1e}", round(residual, 2)  # → ('9.5e-17', 1.28)

In code: residual_chain_gradient multiplies the per-block factors from the table, with or without the skip path.

Why it matters residual connections (ResNet, 2015) made 100+ layer networks trainable, and every transformer wraps both its attention and its feed-forward sub-layers in one (see primer.ml.transformer). One caveat the code shows: adding corrections forever makes the signal grow (a 30-layer residual ReLU stack here grows its gradient about 7 million×), which is why residuals are always paired with normalization.

A teacher can "grade on a curve" in two ways. Batch normalization curves each question across the whole class: your score on question 3 is compared with everyone else's score on question 3, so your grade depends on who else sat the exam. Layer normalization curves each student across their own answers: your scores are rescaled relative to your own average, whoever else is in the room. Both re-centre and re-scale numbers into a steady range so no layer is swamped by huge or tiny values.

Worked examples:

• BatchNorm, batch [[1, 2], [3, 6]]: column means (2, 4), standard deviations (1, 2), so the output is [[−1, −1], [1, 1]]. The value 1 becomes −1.22 in the batch (1, 3, 5) but −0.93 in the batch (1, 3, 11).
• LayerNorm, one row (1, 2, 3, 4): mean 2.5, standard deviation 1.118, so the output is (−1.342, −0.447, 0.447, 1.342), whatever else is in the batch.
• RMSNorm, the same row: root-mean-square √((1+4+9+16)/4) = 2.739, so the output is (0.365, 0.730, 1.095, 1.461). No mean is subtracted.

flowchart LR
  subgraph M["activations: rows = examples, columns = features"]
    direction TB
    r1["ex 1: a  b  c  d"]
    r2["ex 2: e  f  g  h"]
    r3["ex 3: i  j  k  l"]
  end
  M -->|"down each column<br/>(across the batch)"| BN[BatchNorm]
  M -->|"along each row<br/>(within one example)"| LN[LayerNorm / RMSNorm]

BatchNorm is the column version of the same formula, with and computed across the batch for each feature.

In Python:

import math
x = [1, 2, 3, 4]
d, eps, gamma, beta = len(x), 1e-5, 1.0, 0.0
# the row's mean
mu = sum(x) / d
# σ², the row's variance
var = sum((x_i - mu) ** 2 for x_i in x) / d
mu, var  # → (2.5, 1.25)
# LayerNorm
[round(gamma * (x_i - mu) / math.sqrt(var + eps) + beta, 3) for x_i in x]  # → [-1.342, -0.447, 0.447, 1.342]
# √((1/d) Σ x_i² + ε)
rms = math.sqrt(sum(x_i ** 2 for x_i in x) / d + eps)
round(rms, 3)  # → 2.739
# RMSNorm: no mean subtracted
[round(gamma * x_i / rms, 3) for x_i in x]  # → [0.365, 0.73, 1.095, 1.461]

In code: batch_norm normalizes each column across the batch, layer_norm each row across its own features, and rms_norm divides each row by its root-mean-square.

Why it matters transformers use LayerNorm or RMSNorm, never BatchNorm: sequences have different lengths, batches at inference are often size 1, and an example's output must not depend on its batch-mates. Modern LLMs (Llama and others) use RMSNorm because it's cheaper and works as well, and they place it before each sub-layer ("pre-norm"), which keeps the residual path clean and trains more stably. CNNs are where BatchNorm lives on: ResNet puts it after every convolution.

Even with all of the above, one unlucky batch can produce a gradient spike. A circuit breaker doesn't stop the current; it caps it. Clipping by global norm does the same to the update: if the gradient's total length exceeds a limit, it's scaled down to the limit, direction unchanged.

Worked example: in the exploding (too-large) 30-layer network above, the gradients' combined length is astronomically large. Clipped with a limit of 1, the update has length exactly 1 and points the same way.

flowchart LR
  G[Gradients of all layers] --> N["‖g‖: combined length"]
  N --> C{"above the limit?"}
  C -->|yes| S["scale every gradient by limit / ‖g‖"]
  C -->|no| K[leave unchanged]

import math
# a stand-in with the same enormous length
g = [4.8e46, 6.4e46]
# ‖g‖
norm = math.sqrt(sum(g_i ** 2 for g_i in g))
f"{norm:.0e}"  # → '8e+46'
c = 1
# min(1, c / ‖g‖)
scale = min(1, c / norm)
f"{scale:.0e}"  # → '1e-47'
# the new length
round(math.sqrt(sum((g_i * scale) ** 2 for g_i in g)), 6)  # → 1.0

In code: layer_gradients collects every layer's weight gradient from a 30-layer network, and global_norm_after_clipping reports their combined length after clipping.

Why it matters clipping treats the symptom, not the cause: it makes a rare spike harmless, but a network that explodes on every step needs better initialization or normalization. Almost every large training run clips at a norm of about 1.

Answer each one out loud or on paper before you open it. If you can explain it, you know it.