Weight Initialization

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. The Variance Propagation Problem
5. Xavier/Glorot Initialization
6. He/Kaiming Initialization for ReLU
7. Variance in the Backward Pass
8. Orthogonal Initialization
9. Initialization for Residual Networks
10. Common Pitfalls
11. Summary
12. Exercises

Learning Objectives

1. Derive the variance of layer outputs as a function of weights and inputs.
2. Derive the Xavier condition for maintaining variance across layers.
3. Derive the He correction for ReLU networks.
4. Analyze backward-pass variance propagation.
5. Explain orthogonal initialization and its spectral properties.

Notation

• $n_l$ — number of neurons in layer $l$ (fan-in to layer $l+1$)
• $\mathbf{W}^{(l)} \in \mathbb{R}^{n_{l+1} \times n_l}$ — weight matrix at layer $l$
• $\mathbf{h}^{(l)} \in \mathbb{R}^{n_l}$ — activations at layer $l$
• $\text{Var}(h_i^{(l)})$ — variance of activations
• $\text{Var}(w_{ij}^{(l)})$ — variance of weight entries

Core Intuition

If weights are initialized too small, signals shrink exponentially through layers (vanishing activations). If too large, signals explode. Proper initialization maintains the variance of activations (forward pass) and gradients (backward pass) across all layers, ensuring stable training from the first iteration.

The Variance Propagation Problem

Consider layer $l$ with linear pre-activation:

Assuming $w_{ij}$ and $h_j$ are independent, zero-mean, and iid within each layer:

For the signal to neither explode nor vanish, we need:

Xavier/Glorot Initialization

For linear or tanh activations (where $h = \phi(z) \approx z$ near zero):

Forward pass condition: $\text{Var}(w) = 1/n_{\text{in}}$.

Backward pass condition (maintaining gradient variance): $\text{Var}(w) = 1/n_{\text{out}}$.

Xavier compromise (averaging both constraints):

In practice:

• Uniform: $w \sim U\left[-\sqrt{\frac{6}{n_{\text{in}} + n_{\text{out}}}}, \sqrt{\frac{6}{n_{\text{in}} + n_{\text{out}}}}\right]$
• Gaussian: $w \sim \mathcal{N}\left(0, \frac{2}{n_{\text{in}} + n_{\text{out}}}\right)$

He/Kaiming Initialization for ReLU

ReLU sets half the activations to zero: $\mathbb{E}[\text{ReLU}(z)^2] = \frac{1}{2}\text{Var}(z)$ (for zero-mean Gaussian $z$).

Derivation. For $z \sim \mathcal{N}(0, \sigma^2)$:

The variance halves at each ReLU. To compensate:

Setting $\text{Var}(z^{(l)}) = \text{Var}(z^{(l-1)})$:

This is He initialization (He et al., 2015). The factor of 2 compensates for ReLU zeroing half the distribution.

Variance in the Backward Pass

During backpropagation, gradients propagate as:

By the same variance analysis:

For gradient stability: $n_l \cdot \text{Var}(w) = 1$, giving $\text{Var}(w) = 1/n_{\text{out}}$.

Combined: Xavier averages the forward ($1/n_{\text{in}}$) and backward ($1/n_{\text{out}}$) constraints.

Orthogonal Initialization

Idea: Initialize $\mathbf{W}$ as a (scaled) orthogonal matrix.

Construction: Generate a random Gaussian matrix, compute its QR decomposition $\mathbf{G} = \mathbf{Q}\mathbf{R}$, and set $\mathbf{W} = \text{gain} \cdot \mathbf{Q}$.

Property: All singular values of $\mathbf{W}$ equal the gain. This means:

• Forward pass: $\lVert\mathbf{W}\mathbf{h}\rVert = \text{gain} \cdot \lVert\mathbf{h}\rVert$ (exact norm preservation).
• Backward pass: $\lVert\mathbf{W}^T\mathbf{g}\rVert = \text{gain} \cdot \lVert\mathbf{g}\rVert$ (exact gradient norm preservation).

For gain $= 1$: perfect isometry. Prevents both explosion and vanishing for linear networks of arbitrary depth.

Initialization for Residual Networks

In a ResNet with skip connections $\mathbf{h}^{(l+1)} = \mathbf{h}^{(l)} + F(\mathbf{h}^{(l)})$, the variance accumulates:

After $L$ blocks, variance grows linearly. Fixup initialization (Zhang et al., 2019) scales the residual branch by $1/\sqrt{L}$:

so total variance remains $\approx 2\text{Var}(h^{(0)})$ regardless of depth.

Common Pitfalls

Pitfall 1. Using Xavier initialization with ReLU networks. The factor-of-2 correction is essential; without it, activations decay as $(1/2)^L$.

Pitfall 2. Initializing all weights to zero. All neurons compute the same function → gradients are identical → symmetry never breaks. Always use random initialization.

Pitfall 3. Forgetting bias initialization. Biases are typically initialized to zero. For ReLU, a small positive bias (e.g., 0.01) can help prevent initial dead neurons, but is not standard.

Summary

• Xavier: $\text{Var}(w) = 2/(n_{\text{in}} + n_{\text{out}})$ — for sigmoid/tanh.
• He: $\text{Var}(w) = 2/n_{\text{in}}$ — for ReLU (compensates 50% zeroing).
• Orthogonal: exact norm preservation via $\mathbf{W} = \mathbf{Q}$ — strongest guarantee for linear dynamics.
• Residual scaling: $1/\sqrt{L}$ prevents variance growth in deep ResNets.
• Proper initialization ensures stable gradient flow from the very first training step.

Exercises

Exercise 1. Derive equation (2) from the assumptions of independence and zero-mean.

Exercise 2. Show that $\mathbb{E}[\text{ReLU}(z)^2] = \sigma^2/2$ for $z \sim \mathcal{N}(0, \sigma^2)$.

Exercise 3. For a linear network of depth $L$ with Xavier initialization, compute $\text{Var}(h^{(L)})/\text{Var}(h^{(0)})$.

Exercise 4. Prove that an orthogonal matrix preserves norms: $\lVert\mathbf{Q}\mathbf{x}\rVert = \lVert\mathbf{x}\rVert$.

Exercise 5. For a ResNet with $L$ blocks and residual variance $\sigma_F^2$ per block, derive the total output variance and verify the $1/\sqrt{L}$ scaling correction.