Weight Initialization¶

Overview¶

Modern networks self-correct during training — but a bad initialization can stall, explode, or permanently poison a network. Initialization is about keeping variance stable as signals flow forward (activations) and backward (gradients).

Zeroing everything → all neurons compute the same gradient (symmetry breaking never happens).
Xavier/Glorot → var(w) = 2/(fan_in+fan_out): good for tanh/sigmoid symmetric activations.
He/Kaiming → var(w) = 2/fan_in (or 2/fan_out): tuned for ReLU family.
LeCun → var = 1/fan_in (good with LN/self-normalizing nets).

💡 The single most common bug: initializing with torch.randn raw (std 1.0) → activations explode; or with std=1e-4 → gradients vanish. Both are silent and GPU-expensive to find.

Why Variance Matters (the math)¶

For a linear layer y = Wx with bias zeroed:

Var(y) = Var(W) * fan_in * Var(x)

To keep Var(y) = Var(x) through the stack you need Var(W) = 1 / fan_in. Xavier generalizes: Var(W) = 2 / (fan_in + fan_out).

import torch, torch.nn as nn

def check_propagation(n_layers=50, init_fn=None):
    torch.manual_seed(0)
    x = torch.randn(1000)
    for _ in range(n_layers):
        w = torch.empty(1000, 1000)
        if init_fn: init_fn(w)
        else: nn.init.normal_(w, std=0.01)   # "small" but wrong
        x = w @ x
    print(f"std after {n_layers} layers: {x.std().item():.3e}")

check_propagation(init_fn=nn.init.normal_)  # explosion with std=1
# xavier keeps std ~1 across 50 layers:
def xav(w): nn.init.xavier_normal_(w)
check_propagation(init_fn=xav)

The Three Main Recipes¶

1. Xavier (Glorot) — sigmoid/tanh¶

nn.init.xavier_uniform_(layer.weight)
nn.init.xavier_normal_(layer.weight)   # normal vs uniform: picks tails

2. He (Kaiming) — ReLU family¶

nn.init.kaiming_uniform_(layer.weight, mode='fan_in', nonlinearity='relu')
nn.init.kaiming_normal_(layer.weight, a=0.01, nonlinearity='leaky_relu')

a is the negative slope of LeakyReLU — it changes the optimal variance.

3. Orthogonal — RNNs¶

nn.init.orthogonal_(rnn.weight_hh_l0)   # keeps eigenvalues near unit circle

Bias & Residual-Friendly Init¶

def init_linear(m):
    if isinstance(m, nn.Linear):
        nn.init.kaiming_uniform_(m.weight, a=math.sqrt(5))
        if m.bias is not None:
            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(m.weight)
            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
            nn.init.uniform_(m.bias, -bound, bound)

Zero the last-layer bias → if it predicts a distribution, gives a clean "untrained" baseline.
Small std init for output heads (e.g., std=0.02) keeps losses from exploding early.

Common Failure Modes¶

Symptom	Cause	Fix
Loss stuck at log(classes)	zero init / symmetric init	small asymmetric init
Loss = NaN immediately	weights too large	He/Xavier, reduce lr
Hours of flat loss, then works	gradient vanishing	correct var, layer norm
Dead ReLU columns forever	big negative bias init	zero/uniform-small bias

Detecting dead units early¶

def dead_unit_ratio(model, x):
    acts = []
    def hook(m, i, o): acts.append(o)
    h = model[0].register_forward_hook(hook)
    model(x)
    h.remove()
    a = acts[0]
    return (a == 0).float().mean().item()

model = nn.Sequential(nn.Linear(32, 64), nn.ReLU())
print("dead ratio:", dead_unit_ratio(model, torch.randn(8, 32)))

Summary Rules¶

Default init in PyTorch is Kaiming-uniform for most conv/linear — fine as-is for ReLU nets.
ReLU/LReLU → He (fan_in) · tanh/sigmoid → Xavier · RNN → Orthogonal.
Zero-init ONLY for specific last layers (e.g., residual branches: zero_init_gamma).
Always apply a consistent init function:

model.apply(lambda m: (
    nn.init.kaiming_uniform_(m.weight) if isinstance(m, nn.Linear) else None
))

Key Takeaways¶

Init controls variance flow: Var(W)·fan must be ~1 to survive 50+ layers.
Match the init to the activation's nonlinearity, not to taste.
NaN immediately → too large; stuck losses → too small/asymmetric.
Residual branches can start at zero and learn identity first.
PyTorch's defaults are well-tuned — don't "improve" them without measurement.