Saturation from 62.3% to 1.4%: Four Lines of Initialisation, and an Order to Debug In

Part 10 left three defects on the table, and parts 12 and 13 each promised them again: the initial loss should be 3.30 and is 17; far too much of h sits against ±1; gradient loses half its size passing through tanh. This part fixes them.

Four lines change, all of them initialisation. Val goes from 2.1752 to 2.0994 — a drop of 0.076. So the value of this part is not that number. It is that fixing it yields something reusable: given a model that won’t converge, what to check and in what order.

First, keep z and h apart

The rest of this part is about σ, z and h. Confusing those three makes everything below slippery.

z is tanh’s input and can be any real number; h is tanh’s output and always lies between −1 and 1. Below, “z is too large” refers to the input and “h is saturated” to the output being pressed against ±1 — two things, one the cause of the other.

One sample, all the way down

Take the context e m m (indices [5, 13, 13]), with a as the answer. Every parameter comes from randn, that is, standard deviation 1.

Step by step:

① look up   C[5] and C[13] are 10 ordinary reals each, meaningless at initialisation
② concat    x = 3×10 = 30 dimensions   σ_x = 0.867   ✅ healthy
③ @W1       z (1, 200)                 σ_z = 4.916   ❌ amplified 5.7×
④ tanh      h (1, 200)                 σ_h = 0.904   ❌ 61.5% have |h| > 0.99
⑤ @W2       logits (1, 27)             σ = 11.74     ❌ amplified again
⑥ softmax   loss = 45.84               theory: ln 27 = 3.2958

x is healthy for a reason: it is copied straight from C, and C is randn. The illness starts at step three.

Six numbers hold the whole problem

Put the first six entries of z next to what tanh does to them:

z:  2.87     4.30     4.51     5.57     0.40    −7.70
h:  0.9936   0.9996   0.9998   1.0000   0.378  −1.0000
    └────────── all pinned at ±1 ──────────┘   ↑alive  └dead┘

For z = 5.57, h = 1.0000 and the gradient coefficient 1 − h² is 0 — that neuron will never update again. For z = 0.40, h = 0.378 and the coefficient is 0.857; it is healthy.

In a batch, 62.3% of the 200 neurons are in the first state. As usable capacity: 200 × 37.7% ≈ 75 are actually learning, and the other 125 are burning compute.

45.84 and 17 don’t contradict each other

This sample’s loss is 45.84, while part 10 quoted 17. Both are right: 45.84 is a single sample, and an unlucky one — the model happened to pile its probability onto a different character. Averaged over a batch it is 17.

Part 10 also already unpacked that 17: e⁻¹⁷ ≈ 4.1×10⁻⁸, meaning the model assigns the correct answer a probability of roughly one in twenty-four million. It is not guessing; it is confidently wrong. This part won’t repeat that reasoning, only add one thing part 10 didn’t say: a loss below 3.30 is the more suspicious case — that usually means a data leak, or a miscalculated baseline.

An aside: the tuning made the illness worse

Part 10’s network had 6 inputs and 100 hidden neurons, giving σ_z ≈ 2.6. To bring val down, part 12 raised n_embd from 2 to 10 and n_hidden from 100 to 200 — which took the input dimension from 6 to 30.

So the amplification factor went from √6 = 2.4 to √30 = 5.5, σ_z from 2.6 to 4.9, and the saturated fraction from about half to 62.3%. Part 12’s tuning did lower val, and it also made this defect worse. Neither step knew about the other — which is exactly the argument for fixing initialisation before tuning hyperparameters.

Why these three metrics

They are not arbitrary. Each one corresponds to a specific failure mode.

Metric one: the initial loss

L_init = −ln(1/|V|) = ln|V|

27 characters → 3.2958. This is a theoretical value, not an empirical one: a model that has learned nothing should honestly output a uniform distribution.

Check this first: one print gets it, the theoretical value is certain, and it depends on no hyperparameter.

Metric two: the saturated fraction

The target is fewer than 5% with |h| > 0.99. It matters because it decides directly whether gradient can get back through:

∂h/∂z = 1 − h²

A saturated neuron is a dead neuron. The red “saturated” labels under the bars in part 10’s hidden-layer animation are exactly this.

Two stages, kept apart. Saturation late in training is normal — the model really has learned a strong regularity, and confidence is warranted where it is warranted. Saturation at initialisation is pathological: the model has learned nothing, so what is that confidence based on? It was pushed there by random numbers.

Metric three: the per-layer standard deviation ratio

The target is σ_out / σ_in ≈ 1. This is the most fundamental of the three, because error compounds layer by layer:

The first two are symptoms; this is the disease. The initial loss is too large because the ratio at the logits layer is out of control, and saturation happens because the ratio at z is — two symptoms, one cause.

σ_z = √n · σ_W · σ_x

What a single neuron does:

zⱼ = Σᵢ xᵢ · wᵢⱼ, over n = 30 terms

Two statistical rules are enough:

The second one turns on the cross term vanishing: Var(A+B) = E[A²] + 2E[AB] + E[B²], and under independence E[AB] = E[A]E[B] = 0.

Counter-example: Var(a+a) = 4Var(a), not 2 — because a is perfectly correlated with itself and the cross term survives. That counter-example comes back once more later.

Putting them together:

Var(z) = n · σ_x² σ_W²   ⟹   σ_z = √n · σ_W · σ_x

The √n comes from “variances add, standard deviations take the root”. Summing 30 terms amplifies by 5.5×, not 30×, because the signs are mixed and partly cancel. A random walk is the same fact: flip a coin 100 times and you end up about √100 = 10 steps from the origin, not 100.

Checking: √30 × 1.0 × 0.867 = 4.75, measured 4.92 — off by 3.5%. One sample is only 30 numbers, and the formula describes the distribution; over a whole batch (960 numbers) theory says 4.99 and measurement gives 4.96, which is the tight pair.

This section only used the formula. The definitions of expectation and variance, how each of the three rules is derived, and why it is variance rather than standard deviation that adds — the next part is entirely about that.

Solving for σ_W

To make σ_z = σ_x, turn the formula around:

σ_W = 1 / √n

Then add one coefficient. tanh compresses (tanh(1) = 0.76), so pre-amplify a little:

σ_W = gain / √fan_in

This is Kaiming initialisation, proposed by Kaiming He in 2015 — the same line of work that became ResNet. Before it there was Xavier, which did not account for ReLU’s asymmetry.

fan_in, not fan_out

W1.shape        # (30, 200)
                #  ↑    ↑
                # fan_in fan_out

n is how many terms are summed — 30 of them. The 200 only means “this calculation is done 200 times”, and each repetition is independent of the others.

They cannot both hold unless fan_in = fan_out. Xavier splits the difference: √(2/(fan_in+fan_out)).

The output layer is a special case

W2 gets 0.01 rather than Kaiming, because the goal is a different one:

What follows the output layer is softmax, not tanh — it doesn’t mind small, it minds large. With all logits at 0, softmax gives a uniform distribution and the loss is exactly ln 27, which is the honest ignorance from above.

But it cannot be set to 0:

∂L/∂h = ∂L/∂logits · W₂ᵀ = 0

With W2 all zeros nothing gets back, and no earlier layer ever updates. This is the symmetry-breaking problem. Small, but not zero.

A general rule: the last layer is always a special case. Classification shrinks logits, regression matches the target’s mean, generation has its own conventions — every earlier layer is about “don’t distort the signal”, the last one is about “where does it start”.

An honest result: Kaiming loses

Four versions, measured:

The formula loses to a guessed number. The reason is not hard to see: Kaiming optimises for “signal survives many layers”, and this network has two.

The benefit only arrives once the network is deep; the cost (saturation 1.4% → 9.4%) is paid immediately.

Initialisation has no “optimal”, only “matched”. A shallow network can be hand-tuned more aggressively; a deep one has to use the formula or it will not train at all.

The two goals are in genuine conflict: avoiding saturation wants z small, avoiding decay wants z large. Kaiming can only compromise — what actually solves both at once is BatchNorm.

The checklist

Everything above, assembled into “what to do with a new model”.

Step 0: confirm the code is right. 5 samples, 200 steps — does the loss reach ~0? If it doesn’t, stop tuning hyperparameters and go find the bug. Two minutes; it saves hours.

Step 1: check the initial loss.

print(f"initial loss {loss.item():.4f}   theory {math.log(n_classes):.4f}")

Far above the theoretical value → the output layer’s weights are too large → W_last *= 0.01, b_last = 0.

Step 2: check the per-layer activation distribution.

for name, h in activations.items():
    print(f"{name}  σ={h.std():.3f}  saturated={100*(h.abs()>0.99).float().mean():.1f}%")

Look at every layer, not just the first. What matters is the trend in the ratio as depth grows: growing means σ_W is too large, shrinking means too small, stable means move on.

Step 3: check the gradients.

for name, p in model.named_parameters():
    print(f"{name}  grad σ={p.grad.std():.2e}  ratio={p.grad.std()/p.std():.2e}")

The grad/data ratio is the single most useful number, empirically around 1e-3. Far below it means nothing is learning; far above means something is about to blow up.

Step 4: only now, the learning rate.

The order cannot be reversed. Tuning the learning rate before the initialisation is right means finding the best step size from a bad starting point — change the initialisation and the value you found is void. Part 11’s U-shaped sweep was also measured from a random initialisation; with a different one, that whole curve shifts.

Each layer’s problem disguises itself as the next layer’s: bad initialisation → looks like “wrong learning rate”; wrong learning rate → looks like “not enough capacity”; not enough capacity → looks like “not enough data”. Work from the bottom up, not from the symptom.

Where this method stops working

It assumes the statistics of the input are known and stable. In a deep network that assumption breaks:

layer 2’s input = layer 1’s output

Layer 1’s weights keep changing during training, so its output distribution drifts along with them — and everything computed at initialisation stops holding a thousand steps later. This is internal covariate shift.

BatchNorm’s idea is “rather than getting it right, force it” — subtract the mean and divide by the standard deviation at every layer’s output, regardless of what happened upstream. That is the next section.

Four lines changed here, and 0.076 came off. What actually travels out is three things: the initial loss has a theoretical value to check against, the saturated fraction reads out directly whether gradient can flow, and debugging works bottom-up, because every layer’s problem disguises itself as the next layer’s.

Next is a maths catch-up: where σ_z = √n · σ_W · σ_x comes from, and why the metric being watched is variance rather than the mean or the maximum. This part used the formula and moved on; the next one derives it from the definition of expectation.