Xavier and He Initialization

Consider a fully connected layer computing $y = Wx + b$ followed by an activation $a = \phi(y)$, where $x \in \mathbb{R}^{n_{in}}$ is the layer's input, $W \in \mathbb{R}^{n_{out} \times n_{in}}$ is the weight matrix, and $\phi$ is a nonlinearity such as tanh or ReLU. If the entries of $x$ are independent with mean zero and variance $\mathrm{Var}(x)$, and the entries of $W$ are drawn i.i.d. with mean zero and variance $\mathrm{Var}(W)$, independent of $x$, then each output coordinate is a sum of $n_{in}$ independent products, so its variance is

This is the forward variance-propagation equation, and it is the entire engine behind both initialization schemes: it tells you exactly how the variance of the signal scales from one layer to the next as a function of the fan-in $n_{in}$ and the weight variance. If $\mathrm{Var}(W)$ is too large, $\mathrm{Var}(y)$ grows geometrically with depth (exploding activations, and later, exploding gradients). If $\mathrm{Var}(W)$ is too small, $\mathrm{Var}(y)$ shrinks geometrically (vanishing activations, saturating tanh/sigmoid units, dead ReLUs, and vanishing gradients). This is precisely the failure mode described in the The Vanishing and Exploding Gradient Problem: it need not wait for many gradient steps to appear — it is already present in the very first forward pass if initialization is careless.

Xavier / Glorot initialization targets activations that are approximately linear near zero and roughly symmetric (tanh, sigmoid, softsign). Under the approximation $\phi(y) \approx y$ near the origin (valid because weights are small and pre-activations start close to zero), the forward equation says we should set

to keep $\mathrm{Var}(a) \approx \mathrm{Var}(x)$. But Glorot and Bengio also analyzed the backward pass, where the gradient with respect to $x$ satisfies an analogous relation with $n_{out}$ playing the role of fan-in (because the transpose $W^\top$ has $n_{out}$ columns feeding each output gradient coordinate). Requiring simultaneous forward and backward variance preservation gives two competing constraints, $\mathrm{Var}(W) = 1/n_{in}$ and $\mathrm{Var}(W) = 1/n_{out}$; Glorot and Bengio's compromise — the standard "Xavier" formula — is their harmonic-mean-flavored average:

He initialization targets ReLU and its variants (Leaky ReLU, PReLU). ReLU is not symmetric: it maps roughly half of a zero-mean symmetric distribution to exactly zero and passes the other half through unchanged. This halves the variance passed forward at every layer, an effect Xavier's derivation does not account for since it assumes near-linearity. He, Zhang, Ren, and Sun (2015) redo the forward-variance calculation with the ReLU nonlinearity explicitly included and find that to compensate for this factor of two, the weight variance must double:

Both schemes are usually implemented by sampling from a scaled Gaussian or a scaled uniform distribution with the target variance, using either fan-in alone, fan-out alone, or the fan-in/fan-out average, depending on the exact variant. The core terminology to fix: fan-in ($n_{in}$) is the number of input connections to a unit (e.g., for a convolution, $n_{in} = C_{in} \times k_h \times k_w$); fan-out ($n_{out}$) is the number of output connections; and the scheme is chosen according to the activation function used after that weight layer.

The mechanism is a one-time sampling procedure applied to every weight matrix (or convolutional kernel) at network construction time, before any forward or backward pass. Biases are conventionally initialized to zero (or a small constant for ReLU networks, to reduce the chance of dead units at step zero). The procedure below covers Gaussian and uniform variants of both schemes, generalized to convolutional layers. ```text function COMPUTE_FANS(shape, layer_type):

# shape is (out_channels, in_channels, kH, kW) for conv,

# or (n_out, n_in) for a dense layer if layer_type == "dense": n_in, n_out = shape[1], shape[0] else if layer_type == "conv": receptive_field_size = product(shape[2:]) # kH kW (or kDkHkW) n_in = shape[1] receptive_field_size # C_in receptive field n_out = shape[0] receptive_field_size # C_out * receptive field return n_in, n_out

function XAVIER_INIT(shape, layer_type, distribution="uniform"): n_in, n_out = COMPUTE_FANS(shape, layer_type) variance = 2.0 / (n_in + n_out) if distribution == "uniform": limit = sqrt(3.0 * variance) # uniform[-a,a] has var = a^2/3 W = SAMPLE_UNIFORM(-limit, limit, shape) else: # "normal" std = sqrt(variance) W = SAMPLE_NORMAL(mean=0, std=std, shape) return W

function HE_INIT(shape, layer_type, distribution="normal", mode="fan_in"): n_in, n_out = COMPUTE_FANS(shape, layer_type) fan = n_in if mode == "fan_in" else n_out variance = 2.0 / fan if distribution == "normal": std = sqrt(variance) W = SAMPLE_NORMAL(mean=0, std=std, shape) else: # "uniform" limit = sqrt(3.0 * variance) W = SAMPLE_UNIFORM(-limit, limit, shape) return W

function INITIALIZE_NETWORK(layers): for layer in layers: if layer.activation in {"tanh", "sigmoid", "softsign", "linear"}: layer.W = XAVIER_INIT(layer.weight_shape, layer.type) else if layer.activation in {"relu", "leaky_relu", "prelu", "gelu-ish"}: layer.W = HE_INIT(layer.weight_shape, layer.type) layer.b = ZEROS(layer.bias_shape) return layers ``` A few implementation notes that matter in practice: 1. Uniform vs. normal is a free choice; both hit the target variance, they differ only in tail behavior (uniform has bounded support, Gaussian has unbounded but exponentially small tails). Frameworks default to whichever is more convenient, and empirically the difference is negligible. 2. `mode="fan_in"` vs. `"fan_out"` in He initialization: the original paper uses fan-in (preserves forward variance) but preserving backward gradient variance instead uses fan-out; most libraries default to fan-in since forward-pass stability at initialization is the dominant practical concern. 3. Leaky ReLU correction: for Leaky ReLU with negative slope $\alpha$, the exact He-style variance is $\mathrm{Var}(W) = \dfrac{2}{(1+\alpha^2)\,n_{in}}$, since a fraction of the negative half is now scaled by $\alpha$ rather than zeroed, changing the expected squared output. 4. Biases are zero-initialized because the forward-variance argument only concerns the multiplicative weight term; a zero bias does not shift the variance calculation and empirically works well, though small positive bias initialization is sometimes used for ReLU to avoid an initial "all-dead" layer.

Time complexity. Initializing a weight matrix of shape $n_{out} \times n_{in}$ requires exactly one pseudorandom draw per entry, so sampling cost is $O(n_{in} \cdot n_{out})$ per layer, matching the size of the matrix itself — this is asymptotically optimal since every weight must be written at least once. Computing the fan-in/fan-out for a layer is $O(1)$ for dense layers and $O(\text{rank of kernel shape})$ for convolutions (a small constant, since kernel tensors have fixed rank 3 or 4). Across an entire network with $L$ layers and total parameter count $P$, the whole initialization procedure costs $O(P)$ time and $O(P)$ space — the same order as simply allocating the parameters, so initialization is never the bottleneck relative to the $O(P)$ (or worse) cost of a single forward/backward pass.

Correctness argument. The claim to justify is: under stated independence and zero-mean assumptions, the chosen variance keeps $\mathrm{Var}(\text{activation})$ approximately constant across depth, layer after layer. The argument is an induction on depth using the variance-propagation identity. Base case. Assume the network's input $x^{(0)}$ has i.i.d. (or at least uncorrelated, equal-variance) coordinates with variance $\sigma_0^2$ — a property that standard input normalization gives you directly. Inductive step (Xavier, linear-region approximation). Suppose layer $\ell$'s input $x^{(\ell)}$ has variance $\sigma_\ell^2$ per coordinate, weights $W^{(\ell)}$ i.i.d. mean zero, variance $v = 1/n_{in}$, independent of $x^{(\ell)}$. Then

Since $\phi \approx \mathrm{id}$ near 0 for tanh/sigmoid at small pre-activation magnitude, $\mathrm{Var}(x^{(\ell+1)}) = \mathrm{Var}(\phi(y^{(\ell)})) \approx \sigma_\ell^2$. So $\sigma_{\ell+1}^2 \approx \sigma_\ell^2$: the invariant "variance $\approx \sigma_0^2$" is preserved at every layer, by induction, for all depths $L$. Without this correction, variance would evolve as $\sigma_\ell^2 = (n_{in} v)^\ell \sigma_0^2$, which is exponential in depth whenever $n_{in} v \neq 1$ — exactly the exploding/vanishing behavior the scheme eliminates. Inductive step (He, ReLU). Let $y$ be zero-mean symmetric (true if $W$ is zero-mean symmetric and independent of $x$, by a symmetry argument on the joint distribution). For $\phi = \mathrm{ReLU}$,

because exactly (by symmetry) half the probability mass has $y \geq 0$ contributing $y^2$, and the other half has $y<0$ contributing $0$. So $\mathrm{Var}(x^{(\ell+1)}) = \mathbb{E}[\phi(y^{(\ell)})^2] = \tfrac12 \mathrm{Var}(y^{(\ell)}) = \tfrac12 n_{in} v \, \sigma_\ell^2$. Setting $v = 2/n_{in}$ makes this equal $\sigma_\ell^2$ exactly, restoring the invariant. This is precisely why He's constant is Xavier's constant times two: ReLU's rectification is a variance-halving operation that must be compensated multiplicatively.

Honest caveats. The argument relies on assumptions that are only approximately true in real networks: weights and activations are treated as independent (violated once weights are updated by gradient descent, since $W$ and $x$ become statistically coupled through training); the tanh/sigmoid linear approximation breaks down once pre-activations grow large (which is itself controlled, in a virtuous circle, by keeping variance $\approx 1$); and convolutional weight sharing introduces correlations across spatial positions that the i.i.d. fan-in calculation ignores. Additionally, this is purely a forward-pass, initialization-time guarantee — it says nothing about long-run training dynamics, for which normalization layers (BatchNorm, LayerNorm) and residual connections provide complementary, training-time control that these schemes do not replace. Xavier/He initialization and normalization layers are best understood as solving overlapping but distinct parts of the vanishing/exploding problem: initialization fixes the starting point, normalization fixes it throughout training.

Take a 5-layer plain (no normalization, no skip connections) fully connected network, each hidden layer of width $n = 256$, ReLU activations, standard-normal input coordinates with unit variance. We compare three initialization choices at layer construction time and track the empirical variance of the pre-activation at each layer over a batch of $10{,}000$ synthetic input vectors (numbers below are the expected values under the theory; they match simulation to within a few percent).

Reading the table: with the He scheme (middle column), post-activation variance is restored to exactly $1$ after every layer — the invariant from the correctness argument holds at every depth. With Xavier applied to a ReLU network (left column, a common bug when copy-pasting code across activation types), variance shrinks by a factor of $2$ per layer because Xavier's derivation never accounts for ReLU's rectification; after 5 layers activations have lost 97% of their original scale, and gradients backpropagating through this stack will be correspondingly attenuated — the network is not dead, but it will learn very slowly, especially in the earliest layers. With weights $10\times$ too large in variance (right column, an initialization bug rather than a scheme mismatch), variance explodes by roughly $20\times$ per layer, reaching $10^6$ by layer 5; in a real network this manifests as `NaN` losses within the first few forward passes as activations overflow floating point range or saturate any downstream nonlinearity into numerically flat regions with zero gradient. This table is the entire practical lesson in miniature: matching variance to the activation function is not a minor tuning knob, it is the difference between a signal that survives 5 layers intact and one that has decayed by orders of magnitude or blown past numerical range — before a single gradient step has been taken.

The two schemes emerged from a decade in which understanding why deep networks failed to train was as important as any architectural innovation. Xavier Glorot and Yoshua Bengio's 2010 analysis was among the first to give a principled, activation-aware account of initialization variance, replacing folklore heuristics with an explicit derivation grounded in signal-propagation statistics. Kaiming He and coauthors extended this reasoning to the ReLU-dominated architectures that had by then become the default for convolutional networks, and their initialization scheme was published as part of the same paper that pushed plain feedforward ReLU networks to unprecedented depth, directly motivating the search for the residual connections that followed within the same year.

physicist[Kaiming He|born 1984|Researcher whose 2015 work on rectifier-aware initialization and, soon after, residual connections, was central to enabling very deep convolutional networks to train reliably.]

Every major deep learning framework — PyTorch, TensorFlow/Keras, JAX-based libraries — ships Xavier (often named `glorot_uniform`/`glorot_normal`) and He (`he_uniform`/`he_normal`) as default or near-default initializers, selected automatically or by convention based on the layer's declared activation. Convolutional architectures from AlexNet-successors through ResNets, and their descendants in vision, use He initialization as the standard choice for ReLU-family stacks; architectures relying on tanh gating (older recurrent networks, some autoencoder designs) still use Xavier-style initialization for those specific gates. Framework code for building custom layers routinely exposes a `kernel_initializer` argument defaulting to one of these two schemes, meaning that essentially every practitioner training a from-scratch network — whether for a research prototype, a production vision system, or a course assignment — is implicitly relying on this variance analysis to get a usable starting point for gradient descent, before any of the more expensive machinery of Weight Initialization Strategies like pretraining or layerwise adaptive scaling comes into play.

Ask the tutor: How would you derive the correct initialization variance for the GELU or Swish activation function, which are neither purely linear nor a clean half-rectifier like ReLU?