Initialisation, normalisation and residuals

Deep networks did not become trainable because somebody stacked more layers. Three unglamorous components made it possible: initialising weights so the signal neither explodes nor vanishes on the way through, normalising activations so each layer sees inputs in a stable range, and residual connections that give gradients a direct path back. This section covers what each does and why. Get initialisation wrong and a deep network simply does not train, and the failure hands you nothing to work with.

Initialisation#

Weights initialised too large make activations explode through the layers; too small and they vanish. Schemes such as Xavier and He scale the initial weights by the layer's size so the signal keeps a stable magnitude at the start of training. Get this wrong and a deep network does not train at all.

The reasoning is one line of arithmetic. A unit sums n inputs, each multiplied by a weight, so its variance is roughly n times the variance of one weight. Do that fifty times and the signal is multiplied by that factor fifty times over, which either explodes or disappears unless the factor is close to 1. So the weights are drawn with variance about 1 over n, which is Xavier. ReLU then throws away the negative half, halving the variance again, so the correction is 2 over n, which is He.

import numpy as np

def propagate(name, std, depth=50, width=256):
    rng = np.random.default_rng(0)
    h = rng.normal(0, 1, (128, width))
    for i in range(depth):
        W = rng.normal(0, std(width), (width, width))
        h = np.maximum(0, h @ W)               # ReLU network, 50 layers deep
        if i in (0, 9, 49):
            print(name, i + 1, f"{np.sqrt((h ** 2).mean()):.3e}")

propagate("too small (0.01)",  lambda w: 0.01)
propagate("too large (0.1)",   lambda w: 0.1)
propagate("Xavier sqrt(1/n)",  lambda w: (1.0 / w) ** 0.5)
propagate("He     sqrt(2/n)",  lambda w: (2.0 / w) ** 0.5)

#                       layer 1     layer 10    layer 50
# too small (0.01)    1.125e-01   3.555e-10   2.388e-48   nothing reaches the output
# too large (0.1)     1.125e+00   3.555e+00   2.388e+02   and keeps growing
# Xavier sqrt(1/n)    7.033e-01   3.233e-02   1.486e-08   decays: the ReLU half was not accounted for
# He     sqrt(2/n)    9.947e-01   1.035e+00   4.985e-01   stable through all fifty

Read the failure rows carefully, because they describe what you will actually see. Too small gives a forward signal of 1e-48, so every output is effectively zero, the loss sits exactly at its initial value, and nothing moves; it looks like a bug in the data loader. Too large gives activations in the hundreds, then thousands, and the first NaN arrives within a few steps. Neither hands you an error message that mentions initialisation.

Xavier is the instructive row. It is not wrong, it is correct for a symmetric activation such as tanh, and used with ReLU it loses a factor of the square root of 2 per layer, which over fifty layers is the difference between 0.5 and 1e-8. This is why frameworks want to know your activation: PyTorch's kaiming_normal_ takes a nonlinearity argument, and passing the wrong one is a silent, slow failure rather than a loud one.

One case is not about scale at all. If every weight starts at the same value, including zero, then every hidden unit in a layer computes the same thing, receives the same gradient, and is updated identically forever. A thousand-unit layer behaves as one unit, permanently. Randomness is what breaks that symmetry, and it is a separate requirement from getting the magnitude right.

rng = np.random.default_rng(0)
X = rng.normal(0, 1, (32, 4)); y = rng.normal(0, 1, 32)

def first_gradient(W1, W2):
    h = np.maximum(0, X @ W1)
    d_out = (h @ W2 - y[:, None]) / len(y)
    return X.T @ ((d_out @ W2.T) * (h > 0))

same = first_gradient(np.full((4, 6), 0.5), np.full((6, 1), 0.5))
rand = first_gradient(rng.normal(0, 0.5, (4, 6)), rng.normal(0, 0.5, (6, 1)))
print(np.round(same[0], 4))   # [0.0819 0.0819 0.0819 0.0819 0.0819 0.0819]  identical
print(np.round(rand[0], 4))   # [-0.0284 0.0124 0.0156 -0.0488 0.0163 -0.0234]

Normalisation#

Initialisation only fixes the first step. Once the weights start moving, the distribution of activations drifts, and every layer is chasing a moving target set by the layers below it. Normalisation pins that down during training: subtract the mean, divide by the standard deviation, then apply a learned scale and shift so the layer can still choose any range it wants.

def normalise(h, axis, gamma, beta, eps=1e-5):
    mean = h.mean(axis=axis, keepdims=True)
    var = h.var(axis=axis, keepdims=True)
    return gamma * (h - mean) / np.sqrt(var + eps) + beta

h = np.random.default_rng(0).normal(3.0, 7.0, (32, 128))    # [batch, features]

batch_normed = normalise(h, axis=0, gamma=1.0, beta=0.0)    # down the batch, per feature
layer_normed = normalise(h, axis=1, gamma=1.0, beta=0.0)    # across features, per example

print(round(batch_normed.mean(0).max(), 6), round(batch_normed.std(0).min(), 4))   # ~0.0  ~1.0
print(round(layer_normed.mean(1).max(), 6), round(layer_normed.std(1).min(), 4))   # ~0.0  ~1.0

• Batch normalisation rescales activations using statistics across the batch. Effective in vision, and awkward with small batches or sequences.
• Layer normalisation rescales across the features of each example independently. The standard in transformers, since it does not depend on batch composition.

Normalising down the batch means a prediction depends on which other examples happened to be in the batch with it. During training that is a mild regulariser. At inference there is no batch, so batch norm keeps a running average of the statistics it saw and uses those instead, which is a second set of numbers that must be saved with the model and is why model.eval() changes its behaviour. With a batch of size 2 those statistics are noise, and with variable-length sequences the padding contaminates them. Layer norm has neither problem because it never looks outside a single example, which is the entire reason transformers use it.

RMSNorm is the current default in large language models and is layer norm with the mean subtraction removed, keeping only the division by the root mean square. It works about as well and is measurably cheaper, which matters when it runs twice per block across eighty blocks.

Residual connections#

Add a layer's input to its output, so the layer learns a change rather than a whole new representation. This gives gradients a direct path back through the network, which is what made very deep networks trainable. Almost every modern architecture uses them.

The argument is about what the chain rule does to a product. Without a residual, the gradient reaching layer 1 is the product of forty Jacobians, and a product of forty numbers that are each a little below 1 is nothing at all. With a residual, each term becomes 1 plus something, so there is always a path back that multiplies by 1. Forty layers deep, that is the difference between a first layer that trains and one that does not.

import torch, torch.nn as nn

class Net(nn.Module):
    def __init__(self, depth=40, width=128, residual=False):
        super().__init__()
        self.layers = nn.ModuleList([nn.Linear(width, width) for _ in range(depth)])
        self.residual = residual

    def forward(self, x):
        for lin in self.layers:
            out = torch.relu(lin(x))
            x = x + out if self.residual else out       # the entire change
        return x

for residual in (False, True):
    torch.manual_seed(0)
    net = Net(residual=residual)
    net(torch.randn(64, 128)).sum().backward()
    g = [l.weight.grad.norm().item() for l in net.layers]
    print(residual, f"{g[0]:.3e}", f"{g[39]:.3e}", f"{g[39] / g[0]:.3e}")

# residual=False   layer 0: 3.891e-13   layer 39: 2.072e+02   ratio 5.3e+14
# residual=True    layer 0: 6.565e+06   layer 39: 2.796e+07   ratio 4.3e+00

A ratio of 5e14 means the first layer is not training slowly; it is not training. Any learning rate large enough to move it would have destroyed layer 39 long before. With the skip connection the ratio is about 4, so one learning rate suits the whole network, and that is what made depth a parameter you can increase rather than a wall.

class Block(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.ff = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim))

    def forward(self, x):
        return x + self.ff(self.norm(x))     # residual: input plus change