Weight Initialization

Weight initialization is the process of assigning starting values to the learnable parameters of a neural network before training begins. It matters profoundly because poor initialization can cause gradients to vanish or explode across layers, making deep networks effectively untrainable regardless of architecture or learning rate. The counterintuitive fact: initializing all weights to the same value — even a sensible constant like zero — is strictly worse than random noise, because symmetry between neurons means the network can never differentiate what each unit learns.

💡 The Core Insight

The goal of weight initialization is to keep the variance of activations and gradients approximately constant as signals propagate forward and backward through the network. If variance grows layer by layer, activations saturate and gradients explode; if variance shrinks, activations collapse and gradients vanish — either way, learning stalls. The "aha moment" is recognizing that a randomly initialized network is already a carefully engineered system: the scale of random weights must be calibrated to the fan-in and fan-out of each layer so that the network begins in a trainable regime, not by accident but by design.

⚠️ Common Misconception

The most common mistake is treating initialization as a one-time default setting — just use `torch.nn.Linear`'s built-in Kaiming uniform and forget about it. What practitioners get wrong is assuming that a single scheme (e.g., Kaiming He) is universally optimal regardless of activation function or architecture. Kaiming initialization is derived specifically for ReLU networks; applying it to networks with Tanh activations violates its derivation and can cause instability. The correct mental model is that every initialization scheme is a solution to a specific variance-propagation equation tied to a particular activation function — change the activation, and you may need to change the initialization.

How It Works

The Variance Propagation Problem

Consider a single linear layer with no activation function. Given an input vector $\mathbf{x} \in \mathbb{R}^{n_{in}}$ and weight matrix $\mathbf{W} \in \mathbb{R}^{n_{out} \times n_{in}}$ with entries drawn i.i.d. from a distribution with mean 0 and variance $\sigma^2$ , the output $\mathbf{y} = \mathbf{W}\mathbf{x}$ has variance per element:

\quad \text{Var}(y_j) = n_{in} \cdot \sigma^2 \cdot \text{Var}(x_i)

Explanation: Each output neuron accumulates $n_{in}$ independent products $w_{ji} x_i$ , so variance multiplies by the fan-in, meaning the signal amplitude changes by a factor of $\sqrt{n_{in}} \cdot \sigma$ at every layer. Here $n_{in}$ is the number of inputs (fan-in), $\sigma^2$ is the per-weight variance, and $\text{Var}(x_i)$ is the input variance assumed identical across neurons. To keep $\text{Var}(y_j) = \text{Var}(x_i)$ , we need $\sigma^2 = 1/n_{in}$ , which is exactly LeCun initialization.

Xavier / Glorot Initialization

Xavier Glorot and Yoshua Bengio (2010) jointly constrained both the forward pass (preserve activation variance) and the backward pass (preserve gradient variance). The symmetric requirement gives:

\quad \sigma^2 = \frac{2}{n_{in} + n_{out}}

Explanation: By averaging the fan-in and fan-out constraints, Xavier initialization keeps variance stable in both directions simultaneously, which is especially important for symmetric activations like Tanh and Sigmoid. In practice this is often implemented as a uniform distribution:

\quad W_{ij} \sim \mathcal{U}\!\left(-\sqrt{\frac{6}{n_{in}+n_{out}}},\; +\sqrt{\frac{6}{n_{in}+n_{out}}}\right)

Explanation: A uniform distribution $\mathcal{U}(-a, a)$ has variance $a^2/3$ , so setting $a = \sqrt{6/(n_{in}+n_{out})}$ recovers the target $\sigma^2 = 2/(n_{in}+n_{out})$ .

Kaiming / He Initialization

Xavier assumes a linear activation function. ReLU zeroes out half of its inputs, effectively halving the variance at every layer. Kaiming He et al. (2015) corrected for this:

\quad \sigma^2 = \frac{2}{n_{in}}

Explanation: The factor of 2 compensates for the approximately 50% of activations that ReLU kills, restoring the forward-pass variance to its pre-activation level. For the backward pass, the analogous correction uses fan-out: $\sigma^2 = 2/n_{out}$ . PyTorch's default for `nn.Linear` and `nn.Conv2d` is Kaiming uniform with fan-in mode: ```python import torch import torch.nn as nn import math

Manually applying Kaiming uniform (what PyTorch does by default)

layer = nn.Linear(512, 256) nn.init.kaiming_uniform_(layer.weight, a=0, mode='fan_in', nonlinearity='relu') nn.init.zeros_(layer.bias)

Xavier / Glorot uniform for Tanh networks

layer_tanh = nn.Linear(512, 256) nn.init.xavier_uniform_(layer_tanh.weight, gain=nn.init.calculate_gain('tanh')) nn.init.zeros_(layer_tanh.bias)

LeCun normal (recommended for SELU activations)

layer_selu = nn.Linear(512, 256) nn.init.normal_(layer_selu.weight, mean=0.0, std=1.0 / math.sqrt(512)) nn.init.zeros_(layer_selu.bias) ```

The Orthogonal Initialization Alternative

For very deep recurrent or residual networks, orthogonal initialization draws $\mathbf{W}$ from the set of orthogonal matrices (via QR decomposition of a random Gaussian matrix), guaranteeing that:

\quad \mathbf{W}^\top \mathbf{W} = \mathbf{I}

Explanation: An orthogonal weight matrix preserves the L2 norm of any input vector exactly, so neither explosion nor collapse can occur on the forward pass regardless of depth. ```python

Orthogonal initialization for RNNs or deep residual blocks

rnn_layer = nn.Linear(256, 256) nn.init.orthogonal_(rnn_layer.weight, gain=1.0) ```

researcher · Kaiming He | b. 1989 | Kaiming Initialization for ReLU Networks

Kaiming He is a computer vision researcher who completed his PhD at the Chinese University of Hong Kong and later joined Microsoft Research Asia and then Facebook AI Research (FAIR). In his landmark 2015 paper "Delving Deep into Rectifiers," He derived the variance formula that now bears his name, showing empirically that it enabled training of networks with over 30 layers — at a time when training beyond 20 layers was considered practically infeasible. A surprising fact: the same paper also introduced Parametric ReLU (PReLU), but it is the initialization scheme that became the more lasting contribution.

researcher · Xavier Glorot | b. 1986 | Xavier Initialization and Saturating Activation Analysis

Xavier Glorot conducted his foundational work on initialization as a graduate student under Yoshua Bengio at the Université de Montréal. His 2010 paper "Understanding the difficulty of training deep feedforward neural networks" was among the first to systematically measure and explain gradient vanishing as a function of initialization scale, and it introduced what is now called Xavier or Glorot initialization. Remarkably, the paper's secondary finding — that Tanh consistently outperformed Sigmoid because of better gradient flow — helped shift the field away from Sigmoid activations years before ReLU became dominant.

Figure

Variance Propagation Under Different Initialization Schemes

Figure

Symmetry Breaking: Why Zero Initialization Fails

Intuition

🧠 The Mental Model

Imagine you're tuning a large pipe organ where every pipe starts at the exact same length (zero initialization). When you play a chord, every pipe produces the identical sound, and every pipe responds identically to the tuner's adjustments — you can tune forever and the organ remains monotone. Random initialization is like starting each pipe at a slightly different, roughly correct length: now each pipe responds uniquely to tuning, and the full harmonic range becomes reachable. The variance calibration part has an equally clean analogy: imagine passing a whisper through a chain of 20 people. If each person amplifies the message even slightly (say, 1.1×), by person 20 the whisper is a shout that overwhelms all nuance — that's gradient explosion. If each person quiets it slightly (0.9×), by person 20 nothing is audible — that's gradient vanishing. Good initialization is the rule that says "each person should pass the message at exactly the same volume they received it," which is precisely what the $\sigma^2 = 2/n_{in}$ formula enforces. The practical implication: initialization is not about finding the "right" weights (they'll be trained away quickly). It's about finding a starting point where the gradient signal is coherent enough that the optimizer can take its first meaningful step.

Historical Development

⏳ History

1986

Rumelhart, Hinton, and Williams popularize backpropagation in "Learning Representations by Back-propagating Errors," revealing that gradient vanishing in deep networks is a training-time problem, not just an architectural one.

1991

Sepp Hochreiter's diploma thesis provides a rigorous mathematical analysis of vanishing and exploding gradients in recurrent networks, establishing the theoretical foundation that initialization must address.

1998

Yann LeCun et al. publish "Efficient BackProp," recommending weight initialization with variance

1/n_{in}

(LeCun initialization) and input normalization as best practices for shallow networks, becoming the default for over a decade.

2010

Xavier Glorot and Yoshua Bengio publish "Understanding the difficulty of training deep feedforward neural networks," deriving Xavier initialization and providing empirical evidence that initialization scale directly controls gradient flow depth.

2013

Andrew Saxe, James McClelland, and Surya Ganguli prove that orthogonal initialization with linear networks enables exact gradient propagation and dramatically accelerates deep linear network learning.

2015

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun publish "Delving Deep into Rectifiers," deriving He initialization for ReLU networks and demonstrating training of a 30-layer network, establishing Kaiming as the standard for ReLU-based deep learning.

2017

The Fixup initialization paper and related work begin exploring initialization strategies for ResNets that eliminate the need for batch normalization in residual architectures.

2019

Greg Yang and Samuel Schoenholz's "Mean Field Theory of Deep Neural Networks" provides a unified field-theoretic framework predicting the "edge of chaos" where networks are at the boundary between ordered and chaotic signal propagation — now called the "trainability phase diagram."

2022

–present: Large language models like GPT and PaLM adopt scaled initialization schemes (e.g., scaling residual branch weights by

1/\sqrt{N}

where

N

is the number of residual layers) to maintain stable training at billion-parameter scale.

Why It Matters

🔬 Real-World Impact

ResNet-50/101/152 (He et al., 2015)Kaiming He initialization was literally co-developed with ResNets — without it, training 50+ layer residual networks was not reproducible. Every ResNet trained since uses this initialization as its default starting point in PyTorch and TensorFlow.

GPT-2 and GPT-3 (OpenAI, 2019–2020)OpenAI introduced a modified initialization for their transformer residual streams, scaling weights in projection layers by

1/\sqrt{2L}

where

L

is the number of layers — a direct application of orthogonal-style variance control to prevent residual branch variance from accumulating across 96+ layers in GPT-3.

Vision Transformers (ViT, Google, 2020)ViT initializes patch embedding and MLP weights with truncated normal distributions (stddev 0.02), a choice borrowed from BERT that was found empirically to stabilize attention weight scales; incorrect initialization here causes attention entropy collapse on the first forward pass.

Stable Diffusion / DDPM Denoising NetworksU-Net architectures in diffusion models use zero-initialization for the final convolutional layer in each residual block (a technique from the "zero init" convention), so that at the start of training each residual block is an identity function — giving the optimizer a much easier starting landscape.

LSTM and Transformer Training at ScaleHidden-to-hidden recurrent weight matrices are routinely initialized orthogonally to prevent gradient explosion over long sequences — this is the default in many production NLP systems including older LSTM-based translation models at Google.

Neural Architecture Search (NAS)Weight sharing NAS methods like DARTS rely on a single supernet initialization that must generalize across thousands of candidate architectures; poor initialization of the shared weights causes biased architecture rankings and incorrect search results.

Physics-Informed Neural Networks (PINNs)PINNs solving differential equations are extremely sensitive to initialization because the loss landscape has sharp narrow minima; practitioners often use custom initialization derived from the expected frequency content of the solution, extending classical ideas into domain-specific settings.