Weight initialization

How big should the weights be?

You met variance in Phase 2. It measures spread: how far values sit from their average. The standard deviation is its square root. A standard deviation of 0.1 means a variance of 0.01.

A layer's output spread depends on three things:

• $n_{\text{in}}$, the number of inputs each neuron adds up. This is called the fan-in.
• The variance of the weights.
• The variance of the inputs.

In plain words: more inputs means more terms in the sum, so a bigger total. Bigger weights also mean a bigger total. Multiply the three together to get a rough idea of the output's spread.

This rule is a rough guide, not an exact law. It assumes the weights and inputs are random and centered on zero.

Now try it with the numbers from the code above. There are 256 inputs. The weights have standard deviation 1, so variance 1. The inputs have variance 1. The output variance is about 256 × 1 × 1 = 256.

ReLU then turns every negative number into 0. That throws away about half the signal, so roughly 128 is left. Take the square root to turn variance back into spread: √128 ≈ 11.3. So every layer makes the spread about 11 times bigger. That's the elevenfold growth you saw.

To keep the spread steady, you want each layer's multiplier to be 1. That means choosing the weight variance so that $n_{\text{in}}$ × weight variance ÷ 2 = 1. The answer is a weight variance of 2 ÷ 256, or a standard deviation of √(2/256). That's the third setting in the code.

import math
n = 256
std = 1.0
print(round(math.sqrt(n * std**2 / 2), 1))

import math

n_in = 50
for std in [0.5, 0.2]:
    growth = math.sqrt(n_in * std**2 / 2)   # spread multiplier per ReLU layer
    print(f"std {std}: x{growth:.2f} per layer, x{growth**10:,.0f} after 10 layers")

The two standard recipes

That recipe has a name. Two recipes cover most layers. Each is named after the researchers who proposed it. Both use the fan-in, $n_{\text{in}}$. One also uses the fan-out, $n_{\text{out}}$: the number of outputs the layer produces.

In plain words: for a ReLU layer, divide 2 by the number of inputs and take the square root. For a tanh or sigmoid layer, divide 2 by the inputs plus the outputs, then take the square root. Wider layers get smaller weights.

He initialization has the extra factor of 2 because ReLU zeroes about half its inputs. Doubling the weight variance makes up for the half that's lost.

Here's a small example. A layer has 256 inputs and 128 outputs. He gives √(2/256) ≈ 0.0884. Xavier gives √(2/384) ≈ 0.0722. In Python:

import math

n_in, n_out = 256, 128
print("He:    ", round(math.sqrt(2 / n_in), 4))
print("Xavier:", round(math.sqrt(2 / (n_in + n_out)), 4))

In PyTorch and in LLMs

You rarely write these formulas yourself. PyTorch layers start with sensible values already. You can switch to a different recipe in one line:

import torch.nn as nn

layer = nn.Linear(256, 128)            # already initialized: uniform in ±1/√256 for weights and bias
nn.init.kaiming_normal_(layer.weight, nonlinearity="relu")   # switch to He initialization
nn.init.zeros_(layer.bias)

Starting the biases at zero is fine. The random weights already make the neurons differ.

Large language models (LLMs) make their own choices. GPT-2, which is a 2019 language model from OpenAI, draws its weights from a normal distribution with standard deviation 0.02. A few of its layers add their results onto a running total that flows through the whole model. Those layers get scaled down further, by $1/\sqrt{2 \times \text{layers}}$. Here "layers" means the number of blocks the model stacks.

In plain words: the more blocks a model has, the more things get added to that running total. So each one starts smaller, and the total doesn't grow out of control.

A small example: a 12-layer model gives 1/√24 ≈ 0.204. So those layers start with standard deviation 0.02 × 0.204 ≈ 0.0041.

import math

layers = 12
scale = 1 / math.sqrt(2 * layers)
print(round(scale, 3), round(0.02 * scale, 4))

Later lessons cover residual connections and normalization layers, which are two more ways to keep signals in a healthy range. Together with good initialization, they're why networks with 100 layers can train at all.

Symptoms of bad initialization

• The loss is nan or huge from the first step. The weights are probably too large. (It can also be a learning rate problem.)
• The loss barely moves, and the gradients in the early layers are tiny. The weights may be too small. Or sigmoid and tanh units may be stuck at their flat ends, called saturated.
• Several neurons in a layer always have identical values. They started identical.