DL S022 · Weight Initialization — Xavier, He, and μP

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 22 / 80 · Module M04 · ~3 hours

The story · 2010, Bengio's basement, and the wall at depth 10

Here is a story you can reproduce right now, in your terminal, in thirty seconds. Initialize a 10-layer MLP with weights drawn from torch.randn() (mean 0, std 1). Push a random input through it. Print the activation standard deviation at each layer. You will see something like:

layer 0: std 1.00
layer 1: std 22.71
layer 2: std 512.4
layer 3: std 11534
layer 4: std 259812
layer 5: std 5.86e+06
layer 6: std 1.32e+08
layer 7: nan

Or, if you scale the initialization down and use torch.randn() * 0.01:

Either way, the network is unusable. Activations explode or vanish, and gradients follow. In 2009-2010, this was the reason 10-layer nets barely trained. Yoshua Bengio and his student Xavier Glorot were staring at exactly these numbers in the Université de Montréal ML lab, trying to figure out why "deep" learning was so infuriatingly shallow. The community's workaround was pretraining with denoising autoencoders — a laborious layer-by-layer procedure that took weeks and never quite scaled.

Then Glorot published his AISTATS 2010 paper: Understanding the difficulty of training deep feedforward neural networks. The insight is small and elegant. Write down the variance of a linear layer's output in terms of its input variance and the weight-init variance. Solve for the weight variance that makes them equal. Out drops the formula: σ_w = 1/√n_in. Do the same math for the backward pass and you get σ_w = 1/√n_out. Glorot's practical compromise: use the harmonic mean, σ_w = √(2/(n_in + n_out)). This is Xavier / Glorot init. In one line, deep MLPs became trainable end-to-end without pretraining. The pretraining industry died overnight.

Five years later, Kaiming He (then at MSRA, later Meta AI, now MIT) noticed that ReLU — the newly-hot activation function — was a rectifier, not a linear function. It zeros half of a symmetric input. That halves the output variance. Redo the derivation with a factor of 2, and out drops He initialization (2015): σ_w = √(2/n_in). This single tweak was one of the ingredients that made 152-layer ResNets trainable and won ImageNet 2015 at superhuman accuracy.

That is the classical story. But there is a 2024 sequel. Greg Yang and collaborators at Microsoft Research (now xAI) published Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer and a series of follow-ups — the framework known as μ-parameterization (muP). The idea: if you set up your initialization and learning rate and attention scaling in the right coordinated way, the optimal hyperparameters (learning rate, momentum, weight decay, etc.) become width-invariant. Tune a 40M-parameter model, transfer the hyperparameters unchanged to a 6.7B model, converge on the first try. Cerebras, GPT-4, and several 2024-2025 frontier labs use muP internally to skip the millions-of-dollars hyperparameter search that vanilla initialization forces.

Today you are going to (1) derive the variance-preservation condition from four lines of probability, (2) get Xavier and He as two special cases of the same formula, (3) do the "activation-tracking experiment" that makes the whole session click, (4) internalize five special-case initializations that show up in modern models (residual branches, LSTM forget gate, GPT output projection, embedding scaling, muP), and (5) walk out ready to defend your init choice in a code review.

Prerequisites

• Session 005 — variance of a sum of independent RVs. This session's derivation is one line of that.
• Session 008 — why ReLU zeros half the input.
• Session 021 — the "second line of defense" against variance drift.

1 · The variance-preservation goal

Consider one linear layer:

Assume weights W_ij are drawn independently from a distribution with mean 0 and variance σ²_w. Assume the inputs x_j are independent, mean 0, variance σ²_x. Then for one output component:

Because the terms in the sum are independent (W_ij and x_j both random and independent across j), variances add:

(Using E[WX] = E[W]E[X] for independent, mean-0 variables and E[W²] = Var(W) + E[W]² = σ²_w.)

For activation variance to be preserved across the layer — Var(y) = Var(x) — we need:

That is the Xavier init for tanh, derived in four lines. Everything else in this session is a tweak on this equation.

1.1 The backward pass gives a slightly different answer

The exact same derivation applied to the gradient of the loss with respect to x gives:

For gradient variance preservation on the backward pass, we need σ²_w = 1/n_out.

Two conditions, one variable. You cannot satisfy both unless n_in = n_out. Glorot's practical compromise: use the harmonic mean.

2 · Xavier / Glorot initialization

That is Xavier / Glorot init. Approximately preserves both forward activation variance and backward gradient variance. Assumes the activation is roughly linear in its operating range — true for tanh near zero, approximately true for sigmoid.

In PyTorch:

nn.init.xavier_normal_(layer.weight)                  # normal, std = √(2/(n_in+n_out))
nn.init.xavier_uniform_(layer.weight)                 # uniform on ±√(6/(n_in+n_out))

The uniform version has variance (2·bound)² / 12 = bound² / 3 = 2/(n_in+n_out) when bound = √(6/(n_in+n_out)). Same variance, different distribution shape. Uniform is slightly cheaper to sample and (in Glorot's original experiments) marginally better empirically — probably noise.

3 · He / Kaiming initialization — the ReLU tweak

ReLU zeros negative inputs. If the input to ReLU is symmetric around 0 (which it will be if the previous layer output has mean 0), ReLU keeps half of it in expectation. Concretely, if Y ~ Normal(0, σ²), then:

The output variance of a ReLU layer is half its input variance. To compensate — and preserve overall variance through the linear-then-ReLU pair — the weight variance in the linear layer must be doubled:

That is He / Kaiming init. In PyTorch:

nn.init.kaiming_normal_(layer.weight, mode="fan_in", nonlinearity="relu")
nn.init.kaiming_uniform_(layer.weight, mode="fan_in", nonlinearity="relu")

mode="fan_in" uses n_in and preserves forward variance; mode="fan_out" uses n_out and preserves backward variance. fan_in is the standard choice because forward activations blowing up crashes training faster than backward gradients being slightly off.

For GELU and SiLU (the activations in every modern transformer), the exact variance-halving factor is not 2 but close to 2 — approximately 1.7 for GELU. In practice, everyone uses nonlinearity="relu" even for GELU/SiLU networks and it works fine. The 15% variance mismatch is easily absorbed by the LayerNorm two layers downstream.

4 · The activation-tracking experiment (do this once, remember forever)

This is the single experiment that turns "why does init matter?" from vibes into muscle memory.

import torch, torch.nn as nn
 
def check_init(init_fn, n_layers=10, dim=512, use_relu=True):
    x = torch.randn(1, dim)
    print(f"layer -1  std {x.std().item():.4f}  mean {x.mean().item():.4f}")
    for i in range(n_layers):
        layer = nn.Linear(dim, dim, bias=False)
        init_fn(layer.weight)
        x = layer(x)
        if use_relu:
            x = torch.relu(x)
        print(f"layer {i:2d}  std {x.std().item():.4f}  mean {x.mean().item():.4f}")

Try each of:

You will see this exact pattern. It is why the derivation matters — you can predict the pattern by computing "how does variance scale per layer?" and multiplying by depth.

The rule of thumb: if per-layer variance multiplies by factor k, then after L layers variance has grown by k^L. For a 100-layer network, k = 2 gives 10^30 — instant NaN. k = 0.5 gives 10^-30 — instant zero. You need k ≈ 1 to survive depth. Xavier and He each solve for k = 1 under their respective activation assumptions.

5 · Biases — why zero (and one clever exception)

Biases have no variance-preservation role because they add a constant, not a scaled input. Best practice: initialize biases to zero.

The one exception: for the output of a sigmoid/softmax classifier where you know the marginal class distribution, initialize the bias to log(p_i / (1-p_i)) (or log(p_i)) so the network starts by predicting the class prior. Common in object detection (Lin et al. 2017 focal loss paper) where the "background" class has probability ~0.99 and you want the model to start there rather than at uniform 0.5. Speeds up early convergence by a factor of 3-10× on imbalanced classification.

Another useful bias trick: for LSTMs, initialize the forget gate bias to 1 (not 0). This makes the network start out remembering everything, giving gradients a path through time. From Jozefowicz et al. 2015. Every modern LSTM implementation includes this.

6 · Modern init and normalization — a mutual mitigation

BatchNorm and LayerNorm (Session 021) each renormalize activations back toward zero-mean-unit-var. So even if your init is off by a factor of 2, BN/LN corrects it within one layer. This is a large part of why BN made 100-layer networks trainable in 2015 — it turned "get init exactly right" into "just don't be off by 100×."

Does not mean init doesn't matter:

• Before the first normalization layer, init still matters. In a transformer, the embedding layer's output feeds directly into the first attention block's LN. That first LN cannot fix a broken embedding init — an embedding with std 1.0 versus 0.02 gives wildly different behavior even after LN because LN preserves the relative magnitudes within the feature vector.
• Residual connections (ResNets) still need reasonable init on the residual branch — otherwise the residual overwhelms or vanishes relative to the skip.
• Transformers use LN/RMSNorm heavily but still init weights carefully (usually He or a small std like 0.02 for the LM head, following GPT-2).

Rule: use the right init AND the right normalization. Belt and suspenders.

7 · Special-case initializations in modern models

7.1 Residual branches — init to near-zero (Fixup, T-Fixup, ReZero)

In ResNets and transformers, the residual branch (the thing added to the skip) is often initialized so its output is near-zero (e.g., last conv init to zero, or scaled by 1/√L where L is depth). Reason: at the start of training, y = x + residual(x) ≈ x, i.e., the network starts as identity. This makes very deep nets much more trainable.

Three variants:

• Fixup (Zhang et al. 2019): initialize the last conv of each residual branch to zero, scale earlier weights by L^{-1/(2m-2)} where m is layers per block. Trains 10,000-layer ResNets without normalization.
• T-Fixup (Huang et al. 2020): transformer analog. Scales attention/FFN outputs by 1/√L. Trains post-norm transformers without warmup.
• ReZero (Bachlechner et al. 2020): add a learnable scalar α (init to 0) to each residual branch: y = x + α · f(x). Network starts as identity, learns to activate residual branches over training. Simple and effective.

7.2 GPT-style transformers — scaled init on output projection

GPT-2's approach (from Radford et al. 2019):

• Init all weights N(0, 0.02).
• Scale the residual-branch output projection (the last linear in each attention block and each FFN) by 1/√(2N) where N is the number of layers.

The 1/√(2N) prevents residual-branch outputs from summing to a huge value as depth grows. Every GPT-family model — GPT-3, GPT-4, Llama, Gemma, Mistral, DeepSeek — uses some variant of this pattern. Session 040 (nanoGPT rebuild) implements it explicitly.

7.3 Embedding init

Token embeddings in transformers are typically initialized with N(0, 0.02) — the same std as GPT-2's default. Why 0.02 specifically? Historical accident that survives because it works. The typical embedding vector has D=768 to 8192 features, so N(0, 0.02) gives a per-vector norm of ~0.02 · √D = 0.55 to 1.8, which sits comfortably in the range where the following LayerNorm behaves predictably.

7.4 μ-parameterization (muP) — 2022-2025 hyperparameter transfer

Here is the 2020s twist. In vanilla initialization, the optimal learning rate depends on model width. Train a 40M-parameter model at learning rate 3e-4; the optimal for a 6.7B model might be 1.2e-4. Nobody knows without a sweep. Frontier labs used to burn millions of dollars sweeping hyperparameters at scale.

Greg Yang and collaborators showed in Tensor Programs V (Yang et al. 2022) — the "muTransfer" paper — that if you set up initialization AND learning rate AND attention scaling in a coordinated way, the optimal learning rate becomes width-invariant. Tune a 40M model, transfer to 6.7B unchanged, converge on the first try. The specific ingredients:

• Init scales: input weights 1/√fan_in (standard); output weights 1/fan_in (unusual — a 1/fan_in scaling, not 1/√fan_in); hidden weights 1/√fan_in (standard).
• Learning rates: input/output learning rate is a constant; hidden layer learning rate scales as 1/width.
• Attention scaling: instead of dividing attention scores by √d_k, divide by d_k. Called "1/d attention" or "muP attention."

Together, these produce a limit called the "μP limit" where the model behaves consistently across widths. In practice, cloud-scale labs use it to search hyperparameters on a small proxy model.

Public reports on muP wins:

• Cerebras-GPT (Dey et al. 2023): trained 111M to 13B models using muP-transferred hyperparameters, matching Chinchilla optimality.
• MiniCPM (Hu et al. 2024): reports that muP saved ~90% of hyperparameter search cost for their 1B-2B model line.
• The 2025 Nemotron-H technical report describes muP-inspired scaling for their 8B and 56B hybrid Mamba-Transformer models.

muP is an option, not a mandate. Most open-source model families (Llama, Mistral) do not use it because they run enough hyperparameter sweeps that they do not need it. But if you are training at scale from a small compute budget, muP is the state-of-the-art way to save money.

8 · War stories from real production ML code

9 · Diagram — activation variance vs init

And the "predict-before-you-run" mental picture: variance per layer multiplies by k for whichever init/activation combo. After L layers, total variance drift is k^L. You want |log k| < 1/L roughly. For a 100-layer network, that is 1% per layer — very tight tolerance, which is why He plus normalization plus residual connections are all necessary at that depth.

10 · Try it yourself

11 · Recap

Initialization is the first knob and the least glamorous one, but getting it wrong makes deep networks impossible. Derive σ_w = 1/√n_in from variance preservation. Multiply by 2 for ReLU to get He. Take the harmonic mean of fan_in and fan_out for Xavier. Never init weights to zero. Match init to activation. Use normalization as a second layer of defense. Learn muP if you plan to train at scale from a small compute budget. Everything else is convention.

🧠 Retention scaffold

One-line summary (write it in your own words): _______________________________

Spaced review: re-read §1 (derivation) and §4 (activation-tracking experiment) in 24 hours. Revisit the full session on day 7 with focus on §7 (special cases) and §7.4 (muP).

Next session (S023): you have stable init, you have norms, you have a training loop. What learning rate should you actually use? A constant one? Session 023 introduces learning rate schedules — warmup, cosine decay, one-cycle, and the 2024 Warmup-Stable-Decay (WSD) schedule that MiniCPM used to beat Chinchilla-optimal. Getting the schedule right is worth more than a lot of architecture tweaking.

Sticky note (keep on your desk):

• Xavier: √(2/(n_in + n_out)) for tanh/sigmoid.
• He: √(2/n_in) for ReLU-family (GELU, SiLU included in practice).
• Biases: 0 (except LSTM forget-gate = 1, and classifier bias = log-prior for imbalance).
• Weights: never 0.
• Predict variance × k^L across depth L. You want k ≈ 1.

Previous: ← DL S021 · BatchNorm and LayerNorm · Next: DL S023 · Schedulers →