Weight initialization and normalization

The big idea

Every layer multiplies its input by a weight matrix. Do that 20 times, and any small mismatch in scale compounds. If each layer shrinks the signal by 20 percent, after 20 layers it is 0.8^20, about 1 percent of what it was. If each layer grows it by 20 percent, it is 38 times larger. The same happens to gradients on the way back. Either way, deep networks stop learning: signals vanish to nothing or saturate the activations.

Think of a long line of people passing on a message by whispering. If each person speaks a bit more quietly than they heard, the message fades out; if each one speaks a bit louder, it ends as shouting. You want everyone to pass it on at the volume they received it.

Two tools keep the volume steady. Weight initializationChoosing the random starting values of a network's weights, scaled so that activations and gradients keep roughly the same size from layer to layer. sets the starting weights at exactly the right scale for the layer size and activation. NormalizationLayers that rescale activations to a standard mean and spread during training, such as batch norm, layer norm and RMSNorm. layers go further and re-standardize the activations inside the network, at every step of training.

How it works

Step 1: Why the scale depends on the layer width

A neuron with n inputs sums n products w x. If the weights and inputs are independent with mean zero, the variance of the sum is n times the variance of one product:

To keep Var(z) equal to Var(x), you need Var(w) = 1/n. With n = 256 inputs, that is a weight standard deviation of 1/16 = 0.0625. A "reasonable looking" standard deviation of 1.0 makes each layer's sums 16 times wider than its inputs; 0.01 makes them 6 times narrower.

Step 2: Xavier and He initialization

Xavier (Glorot) initializationInitialization with weight variance about 1/n (or 2 / (n_in + n_out)), which keeps signal scale steady through layers with tanh or sigmoid-like activations. Also called Glorot initialization. uses Var(w) = 1/n_in (or 2 / (n_in + n_out) to balance the forward and backward passes). It suits tanh, which is roughly linear near zero.

ReLU zeroes half of its inputs, which halves the variance at every layer. He (Kaiming) initializationInitialization with weight variance 2/n_in, which compensates for ReLU discarding half its inputs and keeps signal scale steady in ReLU networks. Also called Kaiming initialization. doubles the weight variance to compensate: Var(w) = 2/n_in. In the code, a 20-layer ReLU network with Xavier weights ends with activations around 0.001; with He weights it ends near 1.05.

Too large is just as bad: with tanh and weight std 1.0, 87 percent of the units in the first layer are saturated, with outputs beyond plus or minus 0.99, where tanh's slope is almost zero. Biases are usually initialized to zero. Frameworks apply sensible defaults per layer type, but when you build unusual architectures or very deep stacks, this arithmetic is worth checking.

Step 3: Batch normalization

Initialization only fixes the scale at step 0. As weights change during training, activation scales drift again. Batch normalizationA layer that standardizes each feature to zero mean and unit variance using statistics of the current mini-batch, then applies a learned scale and shift. Uses running averages at inference time. (batch norm) fixes it continuously: for each feature, subtract the mean and divide by the standard deviation computed over the current mini-batch, then apply a learned scale gamma and shift beta:

gamma and beta let the network undo the standardization if that helps, so no expressive power is lost. Batch norm made much deeper image networks trainable and allows larger learning rates. Its quirks come from depending on the batch: small batches give noisy statistics, and at inference time there is no batch, so it uses running averages collected during training, which means training and inference behave differently (that is one of the things model.eval() switches).

Step 4: Layer norm and RMSNorm

Layer normalizationA normalization that standardizes each example across its own features, then applies a learned scale and shift. It does not depend on the batch, so it behaves the same in training and inference. standardizes each example across its own features instead of each feature across the batch. Nothing depends on other examples, so it works with batch size 1, with variable-length sequences, and identically in training and inference. That is why transformers and recurrent networks use it.

RMSNormRoot-mean-square normalization: divides each example's features by their root-mean-square and applies a learned scale, without subtracting the mean. A cheaper variant of layer norm. simplifies layer norm further: divide by the root-mean-square of the features, x / sqrt(mean(x^2)), and apply a learned scale, with no mean subtraction and no shift. It is a bit faster and works as well in practice; many recent language models use it.

Step 5: Why normalization stabilizes training

The early explanation for batch norm was reducing "internal covariate shift", each layer's inputs changing as earlier layers learn. The current understanding emphasizes that normalization makes the loss surface smoother and the gradients better behaved, so larger learning rates are safe, and makes the network much less sensitive to the exact initialization. The last line of the depth experiment shows that robustness: weights of std 0.1 are slightly too large for ReLU at width 256 (He would give 0.0884), enough to grow the signal by about 28 percent per layer, or over 100 times across 20 layers. With a layer norm before each activation they stay at a steady 0.54 to 0.6.

Going deeper

Code: scale through depth, and the three norms

import numpy as np

# Push a batch through 20 layers of width 256 and watch the activation scale
def depth_run(init, act, layers=20, n=256, norm=None, seed=0):
    rng = np.random.default_rng(seed)
    h, stds = rng.normal(size=(512, n)), []
    for _ in range(layers):
        W = rng.normal(size=(n, n)) * init(n)
        h = h @ W
        if norm == "layer":                                    # normalize each example over features
            h = (h - h.mean(1, keepdims=True)) / (h.std(1, keepdims=True) + 1e-5)
        h = np.tanh(h) if act == "tanh" else np.maximum(0, h)
        stds.append(h.std())
    return stds

cases = [
    ("tanh, std 0.01", lambda n: 0.01, "tanh", None),
    ("tanh, std 1.0", lambda n: 1.0, "tanh", None),
    ("tanh, Xavier", lambda n: np.sqrt(1 / n), "tanh", None),
    ("ReLU, Xavier", lambda n: np.sqrt(1 / n), "relu", None),
    ("ReLU, He", lambda n: np.sqrt(2 / n), "relu", None),
    ("ReLU, std 0.1 + LayerNorm", lambda n: 0.1, "relu", "layer"),
]
print("activation std after layer:      1        5       10       20")
for name, init, act, norm in cases:
    s = depth_run(init, act, norm=norm)
    print(f"{name:<28}" + "".join(f"{s[i]:9.3g}" for i in (0, 4, 9, 19)))

# tanh with std 1.0: how many units are saturated (|output| > 0.99)?
rng = np.random.default_rng(0)
h = np.tanh(rng.normal(size=(512, 256)) @ rng.normal(size=(256, 256)))
print(f"\ntanh, std 1.0, layer 1: {np.mean(np.abs(h) > 0.99):.0%} of units saturated")

# The three normalizations on one batch: (4 examples, 3 features)
x = np.array([[1.0, 2.0, 3.0], [2.0, 4.0, 6.0], [0.0, 0.0, 9.0], [5.0, 1.0, 0.0]])
bn = (x - x.mean(0)) / x.std(0)                                   # per feature, across the batch
ln = (x - x.mean(1, keepdims=True)) / x.std(1, keepdims=True)    # per example, across features
rms = x / np.sqrt((x ** 2).mean(1, keepdims=True))               # per example, no mean subtraction
print("\nbatch norm (columns now mean 0, std 1):\n", np.round(bn, 2))
print("layer norm (rows now mean 0, std 1):\n", np.round(ln, 2))
print("RMSNorm (rows now root-mean-square 1):\n", np.round(rms, 2))

activation std after layer:      1        5       10       20
tanh, std 0.01                  0.156 0.000102 1.08e-08 1.14e-16
tanh, std 1.0                   0.975    0.974    0.974    0.974
tanh, Xavier                    0.629    0.321     0.23    0.154
ReLU, Xavier                    0.585    0.152   0.0277  0.00103
ReLU, He                        0.827    0.859    0.887     1.05
ReLU, std 0.1 + LayerNorm       0.584    0.565    0.544    0.598

tanh, std 1.0, layer 1: 87% of units saturated

batch norm (columns now mean 0, std 1):
 [[-0.53  0.17 -0.45]
 [ 0.    1.52  0.45]
 [-1.07 -1.18  1.34]
 [ 1.6  -0.51 -1.34]]
layer norm (rows now mean 0, std 1):
 [[-1.22  0.    1.22]
 [-1.22  0.    1.22]
 [-0.71 -0.71  1.41]
 [ 1.39 -0.46 -0.93]]
RMSNorm (rows now root-mean-square 1):
 [[0.46 0.93 1.39]
 [0.46 0.93 1.39]
 [0.   0.   1.73]
 [1.7  0.34 0.  ]]

The table at the bottom shows the axes clearly. After batch norm each column of the 4 by 3 batch has mean 0 and standard deviation 1. After layer norm each row does, so the first two rows, (1, 2, 3) and (2, 4, 6), become identical: layer norm removes each example's overall scale. RMSNorm also maps them to the same row, but keeps the signs and zeros of the original values, because it does not subtract the mean.

Where the norm goes

In transformers, the placement matters. Original transformers normalized after each residual addition (post-norm); most current models normalize the input of each sub-layer instead (pre-norm), which keeps a clean residual path for gradients and trains more stably at depth ( Vanishing and exploding gradients, residual connections ). In convolutional networks, batch norm typically sits between the convolution and the activation.

Other tricks in the same family

• Group norm splits features into groups and normalizes within each, a batch-independent option for small-batch vision training.
• Residual branches are sometimes initialized so that each block starts as nearly the identity (for example, a zero-initialized last layer in each block), so a very deep network starts as a shallow one.
• Weight decay and gradient clipping also keep scales in check during training ( Regularization for deep nets , Vanishing and exploding gradients, residual connections ).

Trade-offs

Common mistakes

• Initializing every weight to zero, so all neurons in a layer get identical gradients and never become different.
• Using a fixed weight scale such as std 1.0 or 0.01 regardless of layer width.
• Forgetting model.eval() with batch norm, so inference uses the statistics of whatever batch it receives.
• Using batch norm with batch sizes of 1 or 2, where the batch statistics are meaningless.
• Applying weight decay to normalization scales and shifts, which usually should be excluded.

In the wild

Convolutional image networks such as ResNets pair He initialization with batch norm after every convolution. Transformers use layer norm or RMSNorm in a pre-norm arrangement, and large language models commonly use RMSNorm. Frameworks initialize linear and convolutional layers with Xavier- or He-style defaults, which is why most practitioners never set initialization by hand, until a custom architecture fails to train and this arithmetic becomes the first thing to check.

Check yourself

Recap