DEV Community

Devanshu Biswas
Devanshu Biswas

Posted on

Why your deep net won't train: weight initialization

I spent an embarrassing amount of time once staring at a loss curve that just would not move. The data was fine. The learning rate was fine. The architecture was a boring stack of linear layers. And still, nothing. The problem turned out to be the most boring thing imaginable: the random numbers I started the weights with.

That is what Day 28 of my DeepLearningFromZero series is about, and I built a little interactive page so you can watch the failure happen in slow motion.

The starting point is not a detail

Training is just gradient descent from wherever you started. Each step nudges the weights a tiny bit. So if your starting weights already push every activation to zero, or off into saturation, or exploding toward infinity, the gradients are either vanishingly small or absurdly large, and the network never digs itself out. For a shallow net you can get away with sloppy init. For a deep one it is the difference between converging in a couple of epochs and never learning at all.

First, don't use zeros

The tempting move is to set all the weights to zero. Clean, symmetric, no bias. It also completely breaks the network.

If every weight in a layer is the same value, every unit in that layer computes the exact same output from the same inputs. Then during backprop every unit gets the exact same gradient, so they all update identically and stay identical forever. A layer with 512 units has the representational power of one unit. Randomness is the thing that breaks this symmetry and lets units specialize. Every real init scheme draws weights at random for exactly this reason.

The one bit of algebra that runs the whole show

Here is the only equation you need. A pre-activation is a weighted sum over fan_in inputs:

z = sum_i  w_i * x_i
Enter fullscreen mode Exit fullscreen mode

If the weights and inputs are independent and zero-mean, the variances add:

Var(z) = fan_in * Var(w) * Var(x)
Enter fullscreen mode Exit fullscreen mode

Look at that fan_in factor. It means each layer multiplies the variance of the signal by fan_in * Var(w). If that product is less than 1, the signal shrinks a little every layer, and after ten layers it has shrunk toward zero. That's vanishing. If it's greater than 1, the signal grows every layer and blows up. That's exploding. Depth compounds it exponentially, which is why deep nets are so much more sensitive to this than shallow ones.

In the demo, pick "Too-small (x0.01)" and you can literally watch the activation standard deviation fall off a cliff on a log axis. By layer 10 it's around 1e-11. The deep layers are seeing nothing.

The fix: keep the variance alive

If the per-layer multiplier is the villain, the fix is obvious. Make it equal to 1. Set Var(w) = 1 / fan_in and the variance of the activations stays put no matter how deep you go.

That's the whole idea behind the two schemes everyone actually uses.

Xavier / Glorot wants both the forward activations and the backward gradients to keep their variance. Forward asks for 1/fan_in, backward asks for 1/fan_out, and you can't have both unless the layer is square, so Glorot splits the difference:

Var(w) = 2 / (fan_in + fan_out)
Enter fullscreen mode Exit fullscreen mode

This assumes an activation that's roughly linear with slope 1 near zero, which is exactly tanh and sigmoid. In the demo, Xavier plus tanh holds the std nearly flat across all ten layers.

He / Kaiming exists because of ReLU. ReLU sets every negative pre-activation to zero, which throws away about half the variance. So you double the forward term to pay it back:

Var(w) = 2 / fan_in
Enter fullscreen mode Exit fullscreen mode

That factor of 2 is the entire difference from Xavier. In the demo, He plus ReLU keeps the std stable to the deepest layer, and if you run Xavier with ReLU instead you can see the signal slowly leak away because it's under-scaled. That leak is the reason He was invented.

The small stuff that still trips people up

A few things worth knowing:

  • Normal and uniform variants of each scheme are interchangeable. They only differ in shape, not variance. A uniform on (-a, a) has variance a^2/3, so the uniform limit is sqrt(3) * std, which is where the 6 in sqrt(6/fan_in) comes from.
  • Biases go to zero. The weights already break symmetry, so there's no reason to randomize them.
  • BatchNorm and LayerNorm re-center activations every layer, so they make a net much less sensitive to init. Less sensitive, not immune. The first few steps, any block without a norm, and very deep residual stacks still care.

In practice you never hand-roll any of this. It's one line: nn.init.kaiming_normal_(w, nonlinearity="relu") for ReLU nets, nn.init.xavier_normal_(w) for tanh/sigmoid and attention projections.

The interactive version lets you flip between all five schemes and both activations and watch the std-vs-depth line and the activation histogram update from a real forward pass. Seeing "Too-small" collapse and "He" stay flat, side by side, made this click for me far better than the equations ever did:

https://dev48v.infy.uk/dl/day28-weight-init.html

Top comments (0)