Deep Learning for Vision

Weight Initialization

How the starting values of a network's weights are chosen, from breaking symmetry to the variance-preserving rules of Glorot and Bengio and of He et al., orthogonal and data-driven schemes, zero-initialized residual branches, and what PyTorch does by default.

intermediate

Weight initialization is the choice of the values a network’s weights take before the first training step. Two requirements shape it. The weights must differ from unit to unit, or the units of a layer stay copies of each other forever; and their scale must suit the layer’s size and activation function, so that the feature maps of an image keep roughly the same magnitude from the first convolution to the last and the gradient keeps its magnitude on the way back. The standard rules, Glorot and Bengio’s (2010) for tanh-like units and He et al.’s (2015) for rectifiers, set the variance of each layer’s weights from its number of inputs or outputs. Later schemes start from orthogonal matrices, measure the variance on real images, or start every residual branch at zero.

Purpose

Backpropagation multiplies the error signal by one factor per layer, and the forward pass multiplies the signal the same way. A vanishing or exploding gradient results when those factors are consistently below or above 1; that article analyses the product of layer Jacobians and the per-layer factor. Initialization is the one moment when the factors can be set directly: before any training, the weights are random numbers whose distribution we choose, so their scale can be picked to make each factor close to 1 on average. Training then starts from a network in which an edge in the input image still produces a measurable response twenty layers up.

When to Use It

Every network trained from scratch needs an initialization, and every framework applies one by default, so the question is whether the default suits the network. It matters most for deep stacks without normalization or shortcuts, where a factor of 0.7 per layer compounds to a factor of about 1,000 over twenty layers. Normalization layers such as batch normalization restore the scale of every channel during the forward pass and make networks much less sensitive to the scale of the weights before them, though not insensitive (see Evidence and Impact). When a pretrained backbone is fine-tuned, only the new layers, typically the classification head, are initialized; the rest start from the pretrained weights (transfer learning).

How It Works

Why not zeros or a constant

If every filter in a convolutional layer starts with the same weights, every output channel computes the same feature map, receives the same gradient and receives the same update, so the channels never become different: the layer behaves like a single channel copied many times. Goodfellow et al. (§8.4) call this need to “break symmetry” perhaps the only property of the initial point known with complete certainty, and note that random initialization from a high-entropy distribution makes it unlikely that any two units compute the same function. All-zero weights are a special case and worse: by our reasoning, with ReLU and zero biases every hidden activation is 0, so the error signals W⊤δ\mathbf{W}^\top \boldsymbol{\delta} passed down are 0 and no hidden weight receives any gradient. Biases, by contrast, can start at zero, which Goodfellow et al. describe as compatible with most weight initializations: the random weights already break the symmetry. For the same reason, by our reasoning, one layer may start at zero while the others are random if its output does not pass through a ReLU at exactly zero (PyTorch’s ReLU has derivative 0 there), as with the classification layer or the last layer of a residual branch, which the residual schemes below exploit.

Keeping the variance through a layer

He et al. (2015), following Glorot and Bengio, track the variance of the responses layer by layer. A unit of layer ll computes yl=wl⊤xl+bly_l = \mathbf{w}_l^\top \mathbf{x}_l + b_l from nln_l inputs. Assume the weights are independent and identically distributed with zero mean, independent of the inputs, which are themselves i.i.d., and set bl=0b_l = 0. Then

Var⁡[yl]=nlVar⁡[wl] E[xl2].\operatorname{Var}[y_l] = n_l \operatorname{Var}[w_l]\, \mathbb{E}[x_l^2].

The last factor is the mean square of the input, not its variance; the two agree only for zero-mean inputs. What it equals depends on the activation xl=f(yl−1)x_l = f(y_{l-1}):

  • Linear or tanh-like units near zero, the case Glorot and Bengio (2010) studied: f′(0)=1f'(0) = 1 and E[xl2]≈Var⁡[yl−1]\mathbb{E}[x_l^2] \approx \operatorname{Var}[y_{l-1}], so the variance is preserved when nlVar⁡[wl]=1n_l \operatorname{Var}[w_l] = 1.
  • ReLU: if yl−1y_{l-1} is symmetric about zero, half of it is zeroed and E[xl2]=12Var⁡[yl−1]\mathbb{E}[x_l^2] = \tfrac{1}{2}\operatorname{Var}[y_{l-1}], so Var⁡[yl]=12nlVar⁡[wl]Var⁡[yl−1]\operatorname{Var}[y_l] = \tfrac{1}{2} n_l \operatorname{Var}[w_l] \operatorname{Var}[y_{l-1}] and the variance is preserved when 12nlVar⁡[wl]=1\tfrac{1}{2} n_l \operatorname{Var}[w_l] = 1.

Over LL layers the per-layer factors multiply, Var⁡[yL]=Var⁡[y1]∏l=2L12nlVar⁡[wl]\operatorname{Var}[y_L] = \operatorname{Var}[y_1] \prod_{l=2}^{L} \tfrac{1}{2} n_l \operatorname{Var}[w_l], which is why He et al. call this product the key to the initialization design: a factor β\beta per layer rescales the signal by βL\beta^L.

The backward pass

The error signal flows through the transposed weights, so in the backward direction a unit’s gradient is a sum over the units it feeds, not over its inputs. For a fully connected layer this count is the layer’s output width; for a convolution with CoutC_\text{out} filters of size k×kk \times k it is n^l=k2Cout\hat{n}_l = k^2 C_\text{out}. Under the same independence assumptions, and with ReLU’s derivative 0 or 1 with equal probability, He et al. obtain

Var⁡[Δxl]=12 n^lVar⁡[wl]Var⁡[Δxl+1],\operatorname{Var}[\Delta x_l] = \tfrac{1}{2}\, \hat{n}_l \operatorname{Var}[w_l] \operatorname{Var}[\Delta x_{l+1}],

where Δx\Delta x is the gradient of the loss with respect to a layer’s input. The backward condition is therefore 12n^lVar⁡[wl]=1\tfrac{1}{2}\hat{n}_l \operatorname{Var}[w_l] = 1, and in the linear case Glorot and Bengio’s is nl+1Var⁡[wl]=1n_{l+1} \operatorname{Var}[w_l] = 1, with nl+1n_{l+1} the size of the next layer.

Glorot and Bengio’s compromise

A layer with different input and output widths cannot satisfy both ninVar⁡[w]=1n_\text{in} \operatorname{Var}[w] = 1 and noutVar⁡[w]=1n_\text{out} \operatorname{Var}[w] = 1. As a compromise, Glorot and Bengio (2010) proposed

Var⁡[w]=2nin+nout,\operatorname{Var}[w] = \frac{2}{n_\text{in} + n_\text{out}},

which meets both conditions when the widths are equal, and implemented it as a uniform distribution. Since U[−a,a]\mathcal{U}[-a, a] has variance a2/3a^2/3, the bound is

w∼U ⁣[−6nin+nout,  6nin+nout],w \sim \mathcal{U}\!\left[-\sqrt{\frac{6}{n_\text{in} + n_\text{out}}},\; \sqrt{\frac{6}{n_\text{in} + n_\text{out}}}\right],

the normalized initialization, widely called Xavier or Glorot initialization (He et al. note that the Caffe library calls it “Xavier”). The heuristic it replaced, U[−1/n,1/n]\mathcal{U}[-1/\sqrt{n}, 1/\sqrt{n}] with nn the size of the previous layer, gives nVar⁡[w]=1/3n \operatorname{Var}[w] = 1/3, a factor that shrinks with every layer. Glorot and Bengio monitored the singular values of each layer’s Jacobian in tanh networks: the average singular value, which for layers of equal width corresponds to the ratio of activation variances between consecutive layers, was about 0.8 with the normalized initialization and about 0.5 with the old heuristic. Their derivation is for fully connected networks.

He et al.’s rectifier initialization

He et al. (2015) carried the derivation over to ReLU networks and to convolutions. The forward condition gives zero-mean Gaussian weights with standard deviation 2/nl\sqrt{2/n_l} and zero biases; the backward condition gives 2/n^l\sqrt{2/\hat{n}_l}. They point out that either condition alone suffices: if the backward products are 1, the forward product becomes ∏l=2Lnl/n^l=c2/dL\prod_{l=2}^{L} n_l/\hat{n}_l = c_2/d_L, the ratio of the second layer’s input channels to the last layer’s filters, which does not shrink exponentially with depth in common designs, and vice versa. They apply the condition to the first layer too, although no ReLU precedes it, because a factor of 2 in one layer does not compound. For a PReLU with initial negative slope aa the condition becomes

12(1+a2) nlVar⁡[wl]=1,\tfrac{1}{2}(1 + a^2)\, n_l \operatorname{Var}[w_l] = 1,

which is the ReLU case for a=0a = 0 and the linear case of Glorot and Bengio for a=1a = 1. The paper’s large ImageNet models were trained from the backward form (their Eqn. 16), and the authors note that both forms made all their models converge. He et al. call the main difference from Glorot and Bengio the treatment of the rectifier; they also note that Glorot and Bengio averaged the forward and backward conditions and used a uniform distribution, and that the conclusion carries over to Gaussians and to either condition alone.

Fan-in and fan-out of a convolution

For a convolution with k×kk \times k kernels, a response depends on k2k^2 positions in each of CinC_\text{in} input channels, and an input value feeds k2k^2 positions in each of CoutC_\text{out} output channels:

nin=k2Cin,nout=k2Cout.n_\text{in} = k^2 C_\text{in}, \qquad n_\text{out} = k^2 C_\text{out}.

He et al. define n=k2cn = k^2 c and n^=k2d\hat{n} = k^2 d this way, with cc input channels and dd filters. PyTorch computes the two numbers from the weight tensor of shape (Cout,Cin/groups,k,k)(C_\text{out}, C_\text{in}/\text{groups}, k, k): fan-in is the second dimension times k2k^2 and fan-out the first dimension times k2k^2. A 3×33 \times 3 convolution from 64 to 128 channels has nin=576n_\text{in} = 576 and nout=1,152n_\text{out} = 1{,}152, so He’s forward rule gives a standard deviation of 2/576≈0.059\sqrt{2/576} \approx 0.059.

Gain

Framework implementations write the standard deviation as gain/fan\text{gain}/\sqrt{\text{fan}} and let the gain carry the activation’s correction. PyTorch’s torch.nn.init.calculate_gain returns 1 for linear, convolution and sigmoid layers, 5/35/3 for tanh, 2\sqrt{2} for ReLU, 2/(1+s2)\sqrt{2/(1 + s^2)} for leaky ReLU with negative slope ss (the PReLU condition above), and 3/43/4 for SELU; its documentation adds that self-normalizing networks need nonlinearity='linear' instead, for a weight variance of 1/N1/N. Only the ReLU and leaky-ReLU values follow from the derivations above; the documentation does not derive the others, and the source code calls the SELU value empirical.

How to Apply It

PyTorch’s functions

The initializers in torch.nn.init (PyTorch 2.14) work in place on a weight tensor:

  • xavier_uniform_(w, gain=1.0) samples U[−a,a]\mathcal{U}[-a, a] with a=gain6/(fanin+fanout)a = \text{gain}\sqrt{6/(\text{fan}_\text{in} + \text{fan}_\text{out})}; xavier_normal_ uses a Gaussian with standard deviation gain2/(fanin+fanout)\text{gain}\sqrt{2/(\text{fan}_\text{in} + \text{fan}_\text{out})}.
  • kaiming_normal_(w, a=0, mode='fan_in', nonlinearity='leaky_relu') uses a Gaussian with standard deviation gain/fan\text{gain}/\sqrt{\text{fan}}; kaiming_uniform_ uses the uniform distribution with the same variance, bound gain3/fan\text{gain}\sqrt{3/\text{fan}}. mode='fan_in' preserves the variance in the forward pass and mode='fan_out' in the backward pass, in the documentation’s words. The default nonlinearity is leaky ReLU with slope a=0, whose gain equals ReLU’s.
  • orthogonal_(w, gain=1) fills the tensor with a (semi-)orthogonal matrix, flattening the trailing dimensions of a convolution kernel.

A common pattern applies one scheme to every layer of a model; in this example (PyTorch 2.14.1) the second convolution, with fan-in 576, receives weights whose standard deviation matches He’s 0.059:

import torch
import torch.nn as nn

torch.manual_seed(0)

def init_weights(m):
    if isinstance(m, (nn.Conv2d, nn.Linear)):
        nn.init.kaiming_normal_(m.weight, mode="fan_in", nonlinearity="relu")
        if m.bias is not None:
            nn.init.zeros_(m.bias)

model = nn.Sequential(nn.Conv2d(3, 64, 3, padding=1), nn.ReLU(), nn.Conv2d(64, 128, 3, padding=1), nn.ReLU())
model.apply(init_weights)   # visits every submodule
print(f"{model[2].weight.std().item():.4f}")   # He: sqrt(2 / (9 * 64)) = 0.0589
0.0590

TorchVision’s ResNets (version 0.29.0 source) initialize convolutions with kaiming_normal_(mode="fan_out", nonlinearity="relu"), the backward form He et al. used, and batch-normalization layers with scale 1 and shift 0.

PyTorch’s default is not He initialization

nn.Linear and nn.Conv2d call kaiming_uniform_(weight, a=math.sqrt(5)) in their reset_parameters method. With a=5a = \sqrt{5} the leaky-ReLU gain is 2/(1+5)=1/3\sqrt{2/(1+5)} = 1/\sqrt{3}, so the bound is 3⋅(1/3)/fanin\sqrt{3} \cdot (1/\sqrt{3})/\sqrt{\text{fan}_\text{in}}:

w∼U ⁣[−1fanin,  1fanin],Var⁡[w]=13 fanin.w \sim \mathcal{U}\!\left[-\frac{1}{\sqrt{\text{fan}_\text{in}}},\; \frac{1}{\sqrt{\text{fan}_\text{in}}}\right], \qquad \operatorname{Var}[w] = \frac{1}{3\,\text{fan}_\text{in}}.

A comment in the source of nn.Linear says as much, and the nn.Conv2d documentation states the same distribution, U(−q,q)\mathcal{U}(-\sqrt{q}, \sqrt{q}) with q=groups/(Cink2)q = \text{groups}/(C_\text{in} k^2) in our notation. Biases are drawn from the same interval. This is the old heuristic that Glorot and Bengio replaced, with faninVar⁡[w]=1/3\text{fan}_\text{in} \operatorname{Var}[w] = 1/3. By our arithmetic, in a ReLU network each layer then multiplies the forward variance by 12⋅13=1/6\tfrac{1}{2} \cdot \tfrac{1}{3} = 1/6, so the activations’ standard deviation falls by about 6≈2.4\sqrt{6} \approx 2.4 per layer until the random biases dominate. In our reading the default is adequate for shallow networks and for layers followed by normalization, while a deep plain ReLU network needs an explicit initialization.

Example: activation statistics and training outcome

The PyTorch script below (run with PyTorch 2.14.1) builds a plain network of twenty 3×33 \times 3 convolutions with 16 channels and ReLU, no shortcuts and no normalization, followed by global average pooling and a linear classifier. The images are synthetic 12×1212 \times 12 pictures of a single horizontal or vertical bar (length 5 to 8 pixels, random position, Gaussian noise of standard deviation 0.2, then standardized). For each initialization it prints the standard deviation of the ReLU feature maps at layers 1, 5, 10, 15 and 20 for 256 training images at initialization, then trains for 200 steps of Adam (learning rate 10−310^{-3}, batches of 32) on 1,024 images and prints the accuracy on 1,024 new ones. The schemes are Gaussian weights of standard deviation 0.01 (“small”), PyTorch’s default, Glorot uniform, He normal with fan-in, LSUV (described under Variants), and the small initialization with batch normalization after every convolution. The script takes the seed as its argument; the output below is for seed 7.

import math, sys
import torch
import torch.nn as nn
import torch.nn.functional as F

torch.set_num_threads(1)
SEED = int(sys.argv[1]) if len(sys.argv) > 1 else 7

def bars(n, size=12, gen=None):
    """Class 0: one horizontal bar, class 1: one vertical bar (length 5-8 px, random position), plus noise."""
    y = torch.randint(0, 2, (n,), generator=gen)
    length = torch.randint(5, 9, (n,), generator=gen)
    a = torch.randint(0, size, (n,), generator=gen)                      # row (or column) of the bar
    b = (torch.rand(n, generator=gen) * (size - length)).long()          # where the bar starts
    x = torch.zeros(n, size, size)
    for i in range(n):
        if y[i] == 0:
            x[i, a[i], b[i]:b[i] + length[i]] = 1
        else:
            x[i, b[i]:b[i] + length[i], a[i]] = 1
    x = x + 0.2 * torch.randn(n, size, size, generator=gen)
    return ((x - x.mean()) / x.std())[:, None], y                       # (N, 1, H, W), standardized pixels

class PlainCNN(nn.Module):
    """20 conv 3x3 layers with 16 channels and ReLU, no shortcuts; global average pooling; linear classifier."""
    def __init__(self, depth=20, C=16, bn=False):
        super().__init__()
        self.convs = nn.ModuleList(nn.Conv2d(1 if l == 0 else C, C, 3, padding=1, bias=not bn) for l in range(depth))
        self.bns = nn.ModuleList(nn.BatchNorm2d(C) if bn else nn.Identity() for _ in range(depth))
        self.head = nn.Linear(C, 2)

    def forward(self, x, stds=None):
        for conv, bn in zip(self.convs, self.bns):
            x = F.relu(bn(conv(x)))
            if stds is not None:
                stds.append(x.std().item())                              # std of this layer's feature maps
        return self.head(x.mean(dim=(2, 3)))

def initialize(model, scheme, xb):
    layers = list(model.convs) + [model.head]
    if scheme == "pytorch default":
        return                                                           # keep kaiming_uniform_(a=sqrt(5)) and its biases
    for m in layers:
        if scheme in ("small", "small + BN"):
            nn.init.normal_(m.weight, 0.0, 0.01)
        elif scheme == "glorot":
            nn.init.xavier_uniform_(m.weight)
        elif scheme == "he":
            nn.init.kaiming_normal_(m.weight, mode="fan_in", nonlinearity="relu")
        elif scheme == "lsuv":
            nn.init.orthogonal_(m.weight)
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    if scheme == "lsuv":                                                 # Mishkin & Matas: scale each layer's output to variance 1
        with torch.no_grad():
            for i, m in enumerate(layers):
                for _ in range(10):
                    h = xb
                    for conv in model.convs[:i]:
                        h = F.relu(conv(h))
                    out = m(h.mean(dim=(2, 3))) if m is model.head else m(h)
                    if abs(out.var().item() - 1.0) < 0.01:
                        break
                    m.weight /= out.var().sqrt()

def run(scheme, steps=200):
    torch.manual_seed(SEED)
    gen = torch.Generator().manual_seed(SEED)
    x_train, y_train = bars(1024, gen=gen)
    x_test, y_test = bars(1024, gen=gen)
    model = PlainCNN(bn=scheme.endswith("BN"))
    initialize(model, scheme, x_train[:256])
    stds = []
    with torch.no_grad():
        model(x_train[:256], stds)                                        # forward statistics at initialization
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    for step in range(steps):
        idx = torch.randint(0, 1024, (32,), generator=gen)
        loss = F.cross_entropy(model(x_train[idx]), y_train[idx])
        opt.zero_grad()
        loss.backward()
        opt.step()
    model.eval()
    with torch.no_grad():
        acc = (model(x_test).argmax(1) == y_test).float().mean().item()
    return stds, acc

print(f"torch {torch.__version__}, seed {SEED}")
print("scheme           std of ReLU output at layer 1      5     10     15     20 | test accuracy")
for scheme in ["small", "pytorch default", "glorot", "he", "lsuv", "small + BN"]:
    stds, acc = run(scheme)
    print(f"{scheme:16s} {stds[0]:26.1e} {stds[4]:6.1e} {stds[9]:6.1e} {stds[14]:6.1e} {stds[19]:6.1e} | {acc:.3f}")

Output:

torch 2.14.1, seed 7
scheme           std of ReLU output at layer 1      5     10     15     20 | test accuracy
small                               1.7e-02 3.7e-07 1.6e-12 3.0e-18 8.6e-24 | 0.501
pytorch default                     3.3e-01 2.7e-02 3.2e-02 2.8e-02 3.8e-02 | 0.501
glorot                              2.1e-01 2.4e-02 3.5e-03 2.2e-04 1.3e-05 | 0.999
he                                  8.2e-01 3.4e-01 3.4e-01 1.4e-01 9.2e-02 | 1.000
lsuv                                6.3e-01 4.7e-01 5.9e-01 7.2e-01 6.7e-01 | 1.000
small + BN                          5.8e-01 5.8e-01 5.8e-01 6.0e-01 6.1e-01 | 0.997

The forward statistics follow the predicted factor 12nVar⁡[w]\sqrt{\tfrac{1}{2} n \operatorname{Var}[w]} per layer (n=144n = 144 after the first layer): about 0.085 for the small initialization and 1/2≈0.71\sqrt{1/2} \approx 0.71 for Glorot’s Var⁡[w]=1/144\operatorname{Var}[w] = 1/144, against measured factors of about 0.08 and 0.6 from layer 5 to layer 20. He initialization and LSUV, which measures the scale instead of deriving it, keep it of order 0.1 to 1. In our reading, the drift left under He initialization comes from the zero padding at the image border, which removes inputs the fan-in count assumes, and from the random fluctuation of a layer with only 16 channels.

Under PyTorch’s default the standard deviation levels off near 0.03 by layer 5. By our arithmetic this is a fixed point set by the default biases: the variance vv of a pre-activation satisfies v=v/6+1/(3⋅144)v = v/6 + 1/(3 \cdot 144), giving a standard deviation of about 0.053 and a ReLU-output standard deviation of about 0.58×0.053≈0.0310.58 \times 0.053 \approx 0.031. At that point the deep feature maps are almost entirely bias; the contribution of the image, shrinking by 1/6≈0.411/\sqrt{6} \approx 0.41 per layer, falls toward the resolution of 32-bit floating point, which in our reading is why this network usually fails to learn. The small initialization fails for another reason, in our reading: its gradients are so small that Adam’s ϵ=10−8\epsilon = 10^{-8} in the denominator of the step dominates, so the updates are negligible. Glorot initialization trains although its top layers start near 10−510^{-5}, in our reading because Adam rescales each parameter’s step by its gradient’s running magnitude (Adam). Batch normalization restores the scale at every layer despite weights of standard deviation 0.01, and that network trains.

The output is for seed 7. Over seeds 0 to 7 (the script’s argument), He initialization, LSUV and the batch-normalized small initialization trained on every seed (test accuracy 0.997–1.000, 0.962–1.000 and 0.915–1.000), Glorot initialization on seven (on seed 2 it stayed at 0.472), PyTorch’s default on two (seeds 0 and 5) and the small initialization on none. Running seed 7 twice gave identical output.

Variants

  • Orthogonal initialization. Saxe, McClelland and Ganguli (2014) analysed deep linear networks and found that scaled Gaussian initializations preserve the norm of a typical vector but not its shape: the product of many Gaussian matrices has most singular values near zero and a few very large ones, so error signals in most directions are attenuated. Random orthogonal matrices, whose product is again orthogonal with all singular values 1, gave learning times on MNIST that did not grow with depth, like greedy layer-wise pre-training. They called the condition that the product of Jacobians act as a near-isometry on as many directions as possible dynamical isometry. For tanh networks with orthogonal weights and a gain gg, they found a critical gain gc=1g_c = 1, the “edge of chaos”, at which a 100-layer, 1,000-unit network kept an O(1)O(1) fraction of the end-to-end Jacobian’s singular values in an O(1)O(1) range. PyTorch’s orthogonal_ implements the scheme for any tensor.
  • Layer-sequential unit-variance (LSUV) initialization. Mishkin and Matas (2016) replaced derivations by measurement: initialize every convolution and fully connected layer with an orthonormal matrix, then, from the first layer to the last, pass a mini-batch through the network and divide the layer’s weights by the standard deviation of its output until the output variance is within a tolerance of 1. The procedure covers activations and layers that, the authors note, the variance formulas did not, such as maxout and max pooling; the authors report that the desired variance was reached in one to five iterations and that tolerances from 0.01 to 0.1 made no noticeable difference. They describe it as orthonormal initialization combined with batch normalization applied only to the first mini-batch. Goodfellow et al. (§8.4) describe the manual version, raising the scale of the first layer with too-small activations until all layers look reasonable on one mini-batch, and credit Mishkin and Matas with specifying it formally.
  • Zero-initialized residual branches. In a residual network, Goyal et al. (2017) set the scale γ\gamma of the last batch-normalization layer in each residual block to 0 instead of 1, so that each block initially computes the identity and the forward and backward signals pass through the shortcuts. TorchVision exposes this as zero_init_residual, off by default as of version 0.29.0.
  • Fixup. Zhang, Dauphin and Ma (2019) derived a lower bound on the gradient norm of a residual network at initialization to explain why, with standard initializations, deep residual networks without normalization cannot be trained at the learning rates normalization allows. Their Fixup initialization sets the classification layer and the last layer of each residual branch to zero, initializes the other layers with a standard method such as He et al.’s and scales the weight layers inside the residual branches by L−1/(2m−2)L^{-1/(2m-2)} for LL branches of mm layers each, and adds a scalar multiplier and scalar biases. The authors identify the rescaling as the essential part.
  • Data-dependent initialization of weight normalization sets each layer’s scale and bias from the statistics of one mini-batch; see normalization layers.

Best Practices

  • Match the rule to the activation. He initialization (gain 2\sqrt{2}) for ReLU, leaky ReLU and PReLU with the (1+a2)(1 + a^2) correction, Glorot initialization for tanh-like and linear layers.
  • Start biases at zero. Goodfellow et al. (§8.4) list exceptions: an output bias set so that the initial predictions match the class frequencies of the training set, and small positive biases, such as 0.1, to keep ReLU units from starting inactive.
  • Keep the output layer small. A classifier whose initial softmax is near uniform starts with a loss near ln⁡K\ln K for KK classes (training a neural network shows the check). He et al. used standard deviations of 0.01 and 0.001 in their fully connected layers, smaller than the derivation gives, because their input pixels ranged over about [−128,128][-128, 128] and larger outputs could overflow the softmax.
  • Look at the statistics. Printing each layer’s activation standard deviation for one mini-batch, as in the example, takes one forward pass and catches a wrong scale before any training.
  • Start residual branches at or near zero in deep residual networks, by a zero γ\gamma in the last normalization layer of each block or, without normalization, by Fixup.

Pitfalls

  • Assuming PyTorch’s default is He initialization. The default bound 1/fanin1/\sqrt{\text{fan}_\text{in}} is a factor of 6\sqrt{6} smaller in standard deviation than He’s uniform bound, and in the example it left a 20-layer ReLU network untrainable on six of eight seeds.
  • A fixed standard deviation for all layers. He et al. computed that if VGG’s “model B”, ten 3×33 \times 3 convolutional layers, were initialized with standard deviation 0.01, it would receive at conv2 a gradient about 1/(1.7×104)1/(1.7 \times 10^4) of what their rule gives, which they suggest explains why some deeper networks initialized with a fixed standard deviation stall. The VGG paper itself prints a variance of 10−210^{-2}, a standard deviation of 0.1. The VGG team, they note, had instead initialized deeper models from a trained model with 8 convolutional layers.
  • Grouped and depthwise convolutions. PyTorch’s fan-in divides the input channels by the number of groups, but its fan-out is k2Coutk^2 C_\text{out} regardless of groups. In a depthwise convolution each input feeds only k2k^2 outputs, so mode='fan_out' overstates the fan-out by a factor equal to the number of groups (our reading of the source); mode='fan_in' is unaffected.
  • Re-initializing after loading pretrained weights. Applying an initialization function to a whole model after loading a checkpoint overwrites the pretrained weights; initialize only the new layers.
  • Unnormalized inputs. The derivations assume inputs of unit-order variance, which raw 8-bit pixels violate (see the output-layer advice above).

Evidence and Impact

He et al. (2015) compared their initialization with Glorot and Bengio’s, implemented as Gaussian weights with nlVar⁡[wl]=1n_l \operatorname{Var}[w_l] = 1, on ReLU networks for ImageNet. A 22-layer model converged with both, theirs reducing the validation error earlier. On a 14-layer model, Glorot’s initialization gave 33.90% top-1 and 13.44% top-5 error and theirs 33.82% and 13.34% (10-view testing on the validation set), so they saw no clear difference in accuracy. A 30-layer model, 27 convolutional and 3 fully connected layers, converged with their initialization, while with Glorot’s the learning stalled completely (their Figure 3) and the gradients, as monitored in their experiments, were diminishing. The 30-layer model reached only 38.56% top-1 and 16.59% top-5 error, worse than the 14-layer model; the authors traced this to higher training error in the deeper model and called the degradation an open problem, the problem the ResNet paper took up.

Glorot and Bengio (2010) reported their normalized initialization’s effect on tanh networks on synthetic Shapeset-3×2 images; the vanishing and exploding gradients article gives the error rates.

Normalization reduces the sensitivity to initialization but does not remove it. Goyal et al. (2017) found that SGD with small mini-batches was not sensitive to initialization because of batch normalization, but that this was not the case for very large mini-batches. On ResNet-50 for ImageNet, zero-initializing the last γ\gamma of each residual block lowered the top-1 validation error from 23.84% to 23.60% with mini-batches of 256 and from 24.11% to 23.74% with mini-batches of 8,192 (means of 5 trials, their Table 2b), a slightly larger gain for the large mini-batch. Without normalization, Zhang et al. (2019) trained residual networks on CIFAR-10 with Fixup at the learning rate used with batch normalization and matched batch normalization’s first-epoch test accuracy up to 10,000 layers, while LSUV and scaling each block’s output by 1/21/\sqrt{2} failed to train at that learning rate beyond 100 layers. On ImageNet their ResNet-50 had 27.6% test error with Fixup against 31.5% for the Xavier-initialized network they cite from Shang et al. and 23.6% with batch normalization; with mixup regularization Fixup reached 24.0%, against 23.9% for group normalization with mixup. Their Table 2 labels these errors “test error”, although its batch-normalization entry is taken from Goyal et al., who report it as top-1 validation error.

Related

  • Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification

    The 2015 ICCV paper by He, Zhang, Ren and Sun of Microsoft Research that introduced the parametric ReLU (PReLU) and the rectifier-aware initialization now called He or Kaiming initialization, and reported 4.94% top-5 error on the ImageNet test set, below the 5.1% reported for a trained human annotator.

  • Training a Neural Network

    The practical workflow for fitting a neural network to labeled images, from splitting the data and normalizing pixels through the training loop, learning curves and early stopping to hyperparameter search, debugging checks, reproducibility and memory.

  • Batch Normalization

    A layer that normalizes each channel of a feature map with the mean and variance of the current mini-batch, then rescales it with learned parameters; it allows higher learning rates in convolutional networks but depends on the batch size and behaves differently in training and inference.

  • Normalization Layers

    Layers that standardize a network's activations with a mean and variance computed over a chosen set of tensor axes, then apply a learned scale and shift; batch, layer, instance, group and RMS normalization differ only in those axes, and weight normalization rescales the weights instead.

  • Residual Connection

    A shortcut that adds a block's input to its output, y = x + F(x), so that the layers learn only a correction to the identity; it lets very deep convolutional networks and transformers train, because signals and gradients always have a path that bypasses the weights.

  • Convolutional Neural Network

    A neural network built from convolutional layers, small learned filters applied at every image position, interleaved with nonlinearities and downsampling, so that features grow from edges to object parts while the parameter count stays independent of image size.

References

  1. Glorot, X. & Bengio, Y. (2010). Understanding the Difficulty of Training Deep Feedforward Neural Networks. Proceedings of the 13th International Conference on Artificial Intelligence and Statistics (AISTATS), PMLR 9, 249–256.
  2. He, K., Zhang, X., Ren, S. & Sun, J. (2015). Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification. IEEE International Conference on Computer Vision (ICCV), 1026–1034.
  3. Saxe, A. M., McClelland, J. L. & Ganguli, S. (2014). Exact Solutions to the Nonlinear Dynamics of Learning in Deep Linear Neural Networks. International Conference on Learning Representations (ICLR). arXiv:1312.6120.
  4. Mishkin, D. & Matas, J. (2016). All You Need Is a Good Init. International Conference on Learning Representations (ICLR). arXiv:1511.06422.
  5. Goyal, P., Dollár, P., Girshick, R., Noordhuis, P., Wesolowski, L., Kyrola, A., Tulloch, A., Jia, Y. & He, K. (2017). Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour. arXiv:1706.02677.
  6. Zhang, H., Dauphin, Y. N. & Ma, T. (2019). Fixup Initialization: Residual Learning Without Normalization. International Conference on Learning Representations (ICLR).
  7. Goodfellow, I., Bengio, Y. & Courville, A. (2016). Deep Learning. MIT Press.
  8. PyTorch documentation (version 2.14): torch.nn.init.
  9. PyTorch documentation (version 2.14): torch.nn.Conv2d.
  10. TorchVision source code (version 0.29.0): torchvision/models/resnet.py.