Initialization

Chapter 1: The Variance Preservation Principle

The explosion/collapse problem from Chapter 0 has a clean mathematical solution. If we can make each layer preserve the variance of its input — variance in equals variance out — then even a 100-layer network will have healthy activations at every layer. The question is: what weight variance achieves this?

The answer turns out to be beautifully simple. And once you see it, the entire zoo of initialization methods — Xavier, He, Lecun — becomes obvious.

Deriving the Variance Formula

Consider a single neuron with n_in input connections. It computes:

We need three assumptions (all standard and approximately true in practice):

• Weights and inputs are independent. This holds because weights are initialized randomly before seeing any data.
• Weights have zero mean. E[w_i] = 0. True for all standard initialization schemes.
• Inputs have zero mean. E[x_i] = 0. Approximately true after centering or batch normalization.

Under these assumptions, the variance of a product of two independent zero-mean random variables is:

And the variance of a sum of independent random variables is:

Since all weights share the same variance Var(w) and all inputs share the same variance Var(x), this simplifies to:

Hand Calculation: The Numbers

Example 1: fan_in = 512, naive initialization N(0, 1)

Weight variance Var(w) = 1. Input variance Var(x) = 1 (standardized input).

Output variance is 512× the input. After 10 layers: 512¹⁰ ≈ 1.3 × 10²⁷. Catastrophic explosion.

Example 2: fan_in = 512, correct initialization Var(w) = 1/512

Output variance equals input variance. After 10 layers: 1¹⁰ = 1. After 100 layers: 1¹⁰⁰ = 1. Perfectly preserved.

Example 3: fan_in = 64

Notice the weights are larger for smaller layers. This makes intuitive sense: with fewer inputs, each weight contributes more to the output, so each individual weight can be larger without blowing up the sum.

See It: Variance at Each Layer

The simulation below shows two panels. Left: variance at each layer for naive N(0,1) initialization (exponential growth). Right: variance for Var(w) = 1/n_in (flat line at 1). Use the fan_in slider to see how layer width affects the explosion rate, and the weight variance slider to find the sweet spot yourself.

Play with this. When fan_in = 512 and Var(w) = 1/512 ≈ 0.00195, the left panel shows a flat line at variance = 1 — perfect preservation. Increase Var(w) even slightly, and the left panel curves upward exponentially. Decrease it, and it curves downward toward zero.

The right panel always shows the correct initialization. Notice how the theoretically correct Var(w) changes with fan_in: wider layers need smaller weights.

From Scratch: Empirical Verification

python
import numpy as np

np.random.seed(42)
fan_in = 512
n_layers = 10
batch_size = 1000

# --- Naive: Var(w) = 1 ---
print("Naive N(0,1):")
x = np.random.randn(batch_size, fan_in)
for i in range(n_layers):
    W = np.random.randn(fan_in, fan_in)  # Var = 1
    x = x @ W
    print(f"  Layer {i}: Var = {np.var(x):.2e}")

# --- Correct: Var(w) = 1/fan_in ---
print(f"\nCorrect Var(w) = 1/{fan_in} = {1/fan_in:.6f}:")
x = np.random.randn(batch_size, fan_in)
for i in range(n_layers):
    W = np.random.randn(fan_in, fan_in) * np.sqrt(1/fan_in)
    x = x @ W
    print(f"  Layer {i}: Var = {np.var(x):.4f}")

Run this. The naive version shows variance doubling at every layer: 500, 260000, 10⁸, ... until overflow. The correct version shows variance hovering near 1.0 at every layer, with small random fluctuations.

The Fan-In Rule

We've arrived at the fundamental principle behind all modern initialization:

This single equation — Var(w) = 1/n_in — is the seed from which Xavier, He, and every other initialization method grows. Xavier extends it to account for the backward pass. He extends it to account for ReLU. But the core idea is always: make each layer preserve variance.

Chapter 2: Xavier/Glorot Initialization

Chapter 1 derived Var(w) = 1/fan_in for the forward pass. But gradients flow backward. For gradients to preserve variance during backpropagation, we need a different condition: Var(w) = 1/fan_out. These are different numbers when the layer isn't square.

Xavier Glorot and Yoshua Bengio noticed this problem in 2010. Their solution is elegant: use the average of the two requirements.

The Backward Pass Problem

During backpropagation, the gradient at layer l flows from layer l+1. The gradient with respect to the input is:

This is another matrix multiplication, but now the matrix is W^T, which has shape fan_out × fan_in. The "effective fan_in" from the gradient's perspective is fan_out. By the same variance argument from Chapter 1:

For gradient variance preservation, we need fan_out · Var(w) = 1, so Var(w) = 1/fan_out.

The Compromise

We have two conflicting requirements:

We can satisfy both simultaneously only when fan_in = fan_out (a square weight matrix). For non-square layers, Xavier uses the harmonic mean — the average of the two:

This doesn't perfectly satisfy either requirement, but it comes close to both. The forward pass variance is slightly off, and the backward pass variance is slightly off, but neither explodes or vanishes.

Two Distributions

Xavier initialization can be implemented with either a normal or uniform distribution:

Xavier Normal:

Xavier Uniform:

Why √6? Because the variance of a uniform distribution U(-a, a) is a²/3. Setting a²/3 = 2/(fan_in + fan_out) gives a² = 6/(fan_in + fan_out), so a = √(6/(fan_in + fan_out)).

Hand Calculation: Concrete Numbers

Layer: fan_in = 512, fan_out = 256.

So W ~ U(-0.0884, 0.0884). Every weight is between -0.09 and +0.09.

Compare with naive N(0, 1): that's weights with std = 1.0 — over 20× too large. It's like setting the volume on a guitar amp to 11 before checking if the mic is plugged in.

Layer: fan_in = fan_out = 512 (square).

When the layer is square, Xavier reduces to exactly the 1/fan_in rule from Chapter 1. The forward and backward requirements agree.

See It: Xavier vs Naive

The simulation below shows a 10-layer tanh network. We show activation histograms at layers 1, 5, and 10. With Xavier, all layers have healthy bell-curve distributions. With naive N(0,1), the later layers saturate — all outputs crowd near ±1, where tanh is flat and gradients vanish.

The naive initialization produces a devastating pattern: at layer 1, the histogram looks fine. By layer 5, values are pushed toward the extremes. By layer 10, every activation is saturated at -1 or +1. The tanh function is completely flat at these extremes, so the gradient is essentially zero. Learning stops.

With Xavier, the histogram at layer 10 looks nearly identical to layer 1 — a smooth bell curve well within the sensitive region of tanh. Gradients flow freely. Learning works.

From Scratch: Implementation

python
import torch
import torch.nn as nn
import numpy as np

# --- Xavier from scratch ---
def xavier_normal(fan_in, fan_out):
    std = np.sqrt(2.0 / (fan_in + fan_out))
    return torch.randn(fan_out, fan_in) * std

def xavier_uniform(fan_in, fan_out):
    a = np.sqrt(6.0 / (fan_in + fan_out))
    return torch.empty(fan_out, fan_in).uniform_(-a, a)

# --- Compare activations through 10 tanh layers ---
torch.manual_seed(42)
fan_in = 512

# Naive
x = torch.randn(64, fan_in)
print("Naive N(0,1) + tanh:")
for i in range(10):
    W = torch.randn(fan_in, fan_in)
    x = torch.tanh(x @ W)
    print(f"  Layer {i}: mean={x.mean():.4f}  std={x.std():.4f}")

# Xavier
torch.manual_seed(42)
x = torch.randn(64, fan_in)
print("\nXavier Normal + tanh:")
for i in range(10):
    W = xavier_normal(fan_in, fan_in)
    x = torch.tanh(x @ W)
    print(f"  Layer {i}: mean={x.mean():.4f}  std={x.std():.4f}")

# PyTorch built-in (identical)
W = torch.empty(512, 512)
nn.init.xavier_uniform_(W)   # in-place initialization
nn.init.xavier_normal_(W)    # normal variant

With naive initialization, you'll see std drop to ~0.05 by layer 5 — the network is saturated and nearly dead. With Xavier, std stays around 0.5-0.65 across all 10 layers. The slight decrease below 1.0 is expected: tanh squishes values, so it always reduces variance slightly. Xavier accounts for this by starting each layer's output at the right scale.

Chapter 3: He/Kaiming Initialization

Xavier assumes the activation function is linear or symmetric. ReLU is neither — it kills everything below zero. If half your neurons output zero, the effective fan_in is halved, and Xavier's variance estimate is too small by a factor of 2. He initialization adds that missing factor.

Published by Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun in 2015 — the same team behind ResNets — this initialization made it possible to train networks with over 100 layers using ReLU activations. Before He init, deep ReLU networks suffered from the slow variance collapse that Xavier couldn't prevent.

The ReLU Problem

Recall from Chapter 1: Var(y) = n_in · Var(w) · Var(x). This assumed linear activations. Now add ReLU:

ReLU(z) = max(0, z). If the pre-activation z follows a symmetric distribution around zero (which it does with zero-mean weights and inputs), then exactly half the values are negative and get zeroed out. The other half pass through unchanged.

What does this do to variance? If we zero out half the values and keep the other half, the expected squared value is halved:

Since E[z] = 0 (zero-mean), E[z²] = Var(z). So:

ReLU halves the variance at every layer. With Xavier's Var(w) = 1/fan_in, the variance after ReLU is:

Each layer halves the variance. After 10 layers: 0.5¹⁰ = 0.001. After 20 layers: 0.5²⁰ = 10^-6. The signal slowly bleeds out. Not as dramatic as N(0,1) explosion, but equally fatal for deep networks.

The Fix: Multiply by 2

The fix is almost comically simple. If ReLU halves the variance, double the weight variance to compensate:

Check it:

The factor of 2 from He exactly cancels the factor of 0.5 from ReLU. Variance is preserved.

Hand Calculation: Xavier vs He

Layer with fan_in = 512, ReLU activation.

With Xavier, each layer loses half its variance. With He, variance is perfectly preserved despite ReLU zeroing out half the neurons.

After 10 ReLU layers:

For Other Activations

He initialization assumes standard ReLU, where exactly half the values are zeroed. For variants:

In practice, He init (factor = 2) works well enough for GELU, SiLU, and Leaky ReLU too. The correction for these activations is between 1.6 and 2.0, so using exactly 2 is close enough.

See It: Xavier vs He with ReLU

The simulation below shows a 10-layer ReLU network. Side-by-side activation histograms at layers 1, 5, and 10. With Xavier, variance dies across layers (histograms narrow progressively). With He, variance is preserved (histograms keep the same spread).

The difference is striking at layer 10. Xavier's histogram has collapsed to a narrow spike near zero — the signal has been halved 10 times, leaving only 0.1% of the original variance. He's histogram at layer 10 looks nearly identical to layer 1. The ReLU is still zeroing half the values (visible as a tall bar at 0), but the surviving values maintain their spread.

From Scratch: Implementation

python
import torch
import torch.nn as nn
import numpy as np

# --- He init from scratch ---
def he_normal(fan_in, fan_out):
    std = np.sqrt(2.0 / fan_in)
    return torch.randn(fan_out, fan_in) * std

def he_uniform(fan_in, fan_out):
    a = np.sqrt(6.0 / fan_in)  # Var = a²/3 = 2/fan_in
    return torch.empty(fan_out, fan_in).uniform_(-a, a)

# --- Compare Xavier vs He through ReLU layers ---
torch.manual_seed(42)
fan_in = 512

# Xavier + ReLU (slow death)
x = torch.randn(64, fan_in)
print("Xavier + ReLU:")
for i in range(10):
    W = torch.randn(fan_in, fan_in) * np.sqrt(1/fan_in)
    x = torch.relu(x @ W)
    print(f"  Layer {i}: mean={x.mean():.4f}  std={x.std():.4f}")

# He + ReLU (stable)
torch.manual_seed(42)
x = torch.randn(64, fan_in)
print("\nHe + ReLU:")
for i in range(10):
    W = he_normal(fan_in, fan_in)
    x = torch.relu(x @ W)
    print(f"  Layer {i}: mean={x.mean():.4f}  std={x.std():.4f}")

# PyTorch built-in
W = torch.empty(512, 512)
nn.init.kaiming_normal_(W, mode='fan_in', nonlinearity='relu')
nn.init.kaiming_uniform_(W, mode='fan_in', nonlinearity='relu')

# mode='fan_in' (default) preserves forward pass variance
# mode='fan_out' preserves backward pass variance
# nonlinearity='relu' uses gain=√2; 'leaky_relu' adjusts for slope

With Xavier + ReLU, std drops to ~0.15 by layer 10 — the network is dying. With He + ReLU, std stays near 0.58 (the expected ReLU output std for unit input) across all 10 layers. The network lives.

The Complete Picture

He initialization is the default for modern deep learning. PyTorch uses Kaiming uniform (He uniform) by default for nn.Linear and nn.Conv2d layers. When you write nn.Linear(512, 256), the weights are automatically initialized with He uniform. You've been using He init all along — you just didn't know it.

Chapter 4: Orthogonal Initialization

Xavier and He use statistical arguments — on average, variance is preserved. Roll the dice a million times and the mean comes out right. But averages are cold comfort when a single unlucky sample explodes layer 47 and kills your training run.

What if we could guarantee norm preservation exactly? Not on average. Not in expectation. For every single input vector, the output has exactly the same magnitude as the input. No stretching, no shrinking, ever.

Such a guarantee exists, and it comes from linear algebra, not statistics. The tool is the orthogonal matrix.

What Is an Orthogonal Matrix?

A square matrix Q is orthogonal if its transpose is its inverse:

This single property has a powerful geometric consequence: multiplying any vector by Q preserves its length. Here's why. Take any vector x and compute the squared norm of Qx:

So ||Qx|| = ||x||. The matrix Q rotates (and possibly reflects) the vector, but never stretches or shrinks it. Think of it like spinning a ball on a table — the ball's size doesn't change, only its orientation.

How to Build One: QR Decomposition

You can't just write down a random orthogonal matrix — random matrices are almost never orthogonal. But you can extract one from any random matrix using QR decomposition.

Every matrix A can be factored as A = QR, where Q is orthogonal and R is upper-triangular. We want Q — we throw away R. The process is:

For rectangular weight matrices (fan_in ≠ fan_out), we generate a larger square matrix, decompose it, and take the appropriate slice of Q. The rows (or columns) of Q are still orthonormal, so the norm-preserving property holds for the dimensions we keep.

Hand Calculation: 2×2 Case

Let's walk through every step with a concrete 2×2 example.

Step 1: Random matrix.

Step 2: QR decomposition. We use Gram-Schmidt. First column of A is a₁ = [0.3, 1.2]. Normalize it:

Second column of A is a₂ = [-0.8, 0.5]. Remove the component along q₁:

Step 3: Our orthogonal matrix Q:

Verification: Q^TQ = I?

Rows are unit-length and perpendicular. That's orthogonality.

Now test norm preservation. Input x = [1, 0], so ||x|| = 1.

The vector rotated from pointing along the x-axis to pointing mostly along the y-axis, but its length is exactly 1. Try any other input — the norm will always be preserved.

The Gain Parameter

Pure orthogonal initialization preserves norms with gain = 1. But sometimes you want controlled scaling. PyTorch's nn.init.orthogonal_ accepts a gain parameter that multiplies Q:

With gain = √2 (appropriate for ReLU), the output norm is √2 times the input norm. This compensates for ReLU zeroing half the activations — exactly like He initialization, but with the rotational structure of an orthogonal matrix.

When to Use Orthogonal Init

Orthogonal initialization shines in three specific scenarios:

Simulation: Rotation vs. Stretching

The widget below shows what happens to a 2D vector when multiplied by a weight matrix. With random init, the vector stretches or shrinks. With orthogonal init, it rotates but keeps the same length. Adjust the gain slider to see controlled scaling. Below: activation variance through 20 layers.

Implementation: From Scratch & PyTorch

python
import numpy as np

def orthogonal_init(shape, gain=1.0):
    """Orthogonal initialization via QR decomposition."""
    # Step 1: Random matrix from N(0, 1)
    rows, cols = shape
    # Make square matrix (larger dimension)
    n = max(rows, cols)
    A = np.random.randn(n, n)

    # Step 2: QR decomposition
    Q, R = np.linalg.qr(A)

    # Fix sign ambiguity (ensures uniform distribution)
    d = np.diag(R)
    Q *= np.sign(d)  # flip columns where R's diagonal is negative

    # Step 3: Slice to desired shape
    Q = Q[:rows, :cols]

    # Step 4: Apply gain
    return gain * Q


# Test: verify norm preservation
W = orthogonal_init((4, 4), gain=1.0)
x = np.random.randn(4)
print(f"||x||  = {np.linalg.norm(x):.4f}")
print(f"||Wx|| = {np.linalg.norm(W @ x):.4f}")
# They will be identical (up to floating-point precision)

python
# PyTorch built-in
import torch.nn as nn

layer = nn.Linear(256, 256)
nn.init.orthogonal_(layer.weight, gain=1.0)
# For ReLU networks, use gain=sqrt(2):
nn.init.orthogonal_(layer.weight, gain=nn.init.calculate_gain('relu'))

Chapter 5: Transformer & LLM Init

Transformers have a structural feature that changes the initialization game entirely: residual connections. Every sublayer — attention and feedforward — adds its output to a running residual stream. And addition accumulates.

Picture a river. Each transformer layer is a tributary that pours into it. With 96 layers (GPT-3), that's 192 tributaries (attention + FFN each). Even if each tributary carries a modest flow, 192 of them together create a flood.

The Residual Accumulation Problem

In a standard transformer block, the computation is:

After one block, x contains the original signal plus two additive terms. After N blocks, x has accumulated 2N additive terms:

If the original x₀ has variance 1, and each sublayer output f_i(x) has variance σ², then the total variance of x_final is approximately:

We want this to stay near 1. Solving for σ²:

Each sublayer's output projection must be initialized with a standard deviation scaled down by 1/√(2N). More layers → smaller initial contributions. This is the core insight behind every modern LLM initialization recipe.

The GPT-2 Recipe

GPT-2's initialization code (and GPT-3, which follows the same pattern) uses a base standard deviation of 0.02 for most weights. But the output projections in attention and FFN get a special scaling:

This 0.02 base isn't arbitrary. For GPT-2's d_model = 768:

GPT-2 uses 0.02, which is conservatively smaller than 0.036. This extra conservatism gives a margin of safety — it's better to start too small (slow initial learning) than too large (unstable training).

Let's compute the residual scaling for different GPT-2 variants:

The pattern is clear: deeper models use dramatically smaller initialization for their residual projections. A 48-layer model's residual weights start 5× smaller than a 12-layer model's.

Hand Calculation: 48-Layer Transformer

Let's verify the math for GPT-2 XL (48 layers). Does residual scaling actually keep the variance stable?

Without scaling:

• Each sublayer output has variance σ² = 0.02² = 0.0004
• After 96 sublayers (48 attention + 48 FFN): accumulated variance = 96 × 0.0004 = 0.0384
• Total variance = 1 + 0.0384 = 1.038 — this doesn't look bad!

Wait — that's because 0.02 is already small. The real problem is what happens inside each sublayer. The output projection multiplies by a d_model × d_model matrix. With standard init (σ = 0.02 per element), the output variance is approximately d_model × 0.02² × input_variance. For d_model = 1600:

• After 96 sublayers: accumulated variance = 96 × 0.64 = 61.4
• Total variance = 1 + 61.4 = 62.4 — the residual stream has 62× the variance it should!

With scaling:

• Residual σ = 0.02 / √96 = 0.00204
• Per-element variance = 0.00204² = 4.16 × 10^-6
• Sublayer output variance ≈ 1600 × 4.16 × 10^-6 = 0.00666
• After 96 sublayers: accumulated variance = 96 × 0.00666 = 0.639
• Total variance = 1 + 0.639 = 1.64 — much more stable!

μP (Maximal Update Parametrization)

The GPT-2 recipe works but has a limitation: when you change model size, you need to re-tune learning rate and other hyperparameters. μP (Maximal Update Parametrization) from Yang & Hu (2021) fixes this by adjusting initialization, learning rate, and activation scaling together in a principled way.

The key insight: μP sets init and LR so that the size of parameter updates relative to parameter magnitudes stays constant across model widths. You tune hyperparameters on a small model and they transfer directly to the large model. Cerebras and some research labs use μP for efficient hyperparameter search.

Simulation: Residual Stream Variance

Adjust the number of layers and toggle residual scaling on/off. Watch how variance grows through the network.

Implementation: GPT-2 Style Init

python
import torch
import torch.nn as nn

class GPT2Block(nn.Module):
    def __init__(self, d_model, n_heads, n_layers):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = nn.MultiheadAttention(d_model, n_heads)
        self.ln2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, 4 * d_model),
            nn.GELU(),
            nn.Linear(4 * d_model, d_model),  # output projection
        )
        self.n_layers = n_layers

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            # Standard init: N(0, 0.02)
            torch.nn.init.normal_(module.weight, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)

        # Special: scale down residual projections
        # These are the LAST linear layers in attention & FFN
        if hasattr(module, '_is_residual_proj'):
            std = 0.02 / ((2 * self.n_layers) ** 0.5)
            torch.nn.init.normal_(module.weight, std=std)

# For GPT-2 XL (48 layers, d_model=1600):
# Standard layers: std = 0.02
# Residual projections: std = 0.02 / sqrt(96) = 0.00204

Chapter 6: The Deep Network Explorer

Time to put everything together. This is the showcase simulation — a full deep network where you pick the initialization method, activation function, number of layers, and layer width, then watch what happens to activations and gradients as signals propagate through the network.

You've learned the theory: Xavier for tanh, He for ReLU, orthogonal for exact preservation, residual scaling for transformers. Now break things. Pick the wrong init for the wrong activation. Make the network absurdly deep. Watch activations collapse to zero or explode to infinity. Then fix it with the right initialization and see everything turn green.

Experiments to Try

Here's what to look for with each combination. Try these in order:

No quiz for the showcase chapter — the simulation is the test. If you can predict what will happen before clicking each combination, you've mastered initialization.

Chapter 7: Zero Init, Bias Init & Practical Recipes

We've covered the general strategies: Xavier for tanh/sigmoid, He for ReLU, orthogonal for exact preservation, residual scaling for transformers. Now the specific tricks. What do you initialize to zero? What do you initialize to a small constant? And what does the actual init code in Llama look like?

Zero Init for Residual Branches

Here's a counterintuitive idea: initialize the last layer of each residual block to all zeros.

In a residual block, the computation is:

If f's last layer has weights W = 0, then f(x) = 0, so:

The network starts as the identity function. Every residual block passes its input straight through, unchanged. The network then gradually learns to deviate from identity — each block making small, incremental modifications rather than random large ones.

This technique is used in:

• GPT-2/GPT-3: Output projections of attention and FFN are scaled down (not pure zero, but toward zero)
• ResNet-v2 (Pre-activation ResNet): BatchNorm's γ parameter initialized to 0 for the last norm in each block
• FixUp initialization: Explicitly zeros the last layer of each residual branch, allowing training without normalization layers

Bias Initialization

Biases are simpler than weights. The standard rule: initialize all biases to zero. Here's why this works:

A bias adds a constant to every neuron's pre-activation: z = Wx + b. If the weights are properly initialized (Xavier, He, etc.), the pre-activations already have zero mean. A zero bias preserves this property. A nonzero bias would shift all neurons in the same direction, which the network would need to unlearn.

There's one historical exception. With ReLU activation, some practitioners initialized biases to a small positive value like 0.01:

The argument: if a neuron's pre-activation is slightly negative at init, ReLU kills it (output = 0). A small positive bias shifts the distribution so more neurons fire. In practice, this made little difference when combined with proper He initialization, and modern networks use zero biases universally.

One more case: many modern LLMs have no biases at all. Llama, PaLM, and other recent architectures remove biases entirely from linear layers. The argument: biases add parameters but contribute little when combined with layer normalization, which re-centers activations anyway.

The Complete LLM Recipe

Let's look at what modern LLMs actually do. Here's the initialization recipe for three major architectures:

The pattern across all three: (1) most weights use a small normal distribution, (2) normalization layers start as identity (γ=1, β=0), (3) output projections are scaled down for deep models.

Hand Calculation: Llama's Init

Llama 2 7B has d_model = 4096 and 32 layers. Let's compute every init value:

Embeddings:

Linear layers (Q, K, V, gate, up projections):

He normal: σ = √(2/fan_in). For a 4096 → 4096 layer:

Output projections (attn_out, ffn_down):

RMSNorm weights:

That's the entire init. No biases (Llama removes them). No β (RMSNorm doesn't have one). Clean and minimal.

The Symmetry Breaking Problem

Let's verify this with a quick hand calculation. Take a 2-layer network with 2 neurons each, all weights initialized to 0.5, bias 0:

Forward pass with input [1, 1]:

• Neuron 1: 0.5(1) + 0.5(1) = 1.0
• Neuron 2: 0.5(1) + 0.5(1) = 1.0

Both neurons produce identical output: 1.0.

Backward pass:

• Both neurons receive identical gradients (because their outputs were identical and they have identical weights to the next layer)
• Both neurons update identically: w_new = 0.5 - η × (same gradient)

After the update, the weights are still identical. Repeat forever — the two neurons are permanently locked in sync. You paid for two neurons but only got one.

Simulation: Recipe Builder

Select a component and an architecture to see its initialization value. Toggle between GPT-2, Llama, and BERT recipes.

Implementation: Complete LLM Init

python
import torch
import torch.nn as nn
import math

def init_llama(model, d_model, n_layers):
    """Initialize a Llama-style model."""
    for name, param in model.named_parameters():
        if param.dim() < 2:
            continue  # skip 1D params (norms)

        if 'embed' in name:
            # Embeddings: N(0, 1/sqrt(d_model))
            nn.init.normal_(param, std=1.0 / math.sqrt(d_model))

        elif 'out_proj' in name or 'down_proj' in name:
            # Residual projections: He normal / sqrt(2*n_layers)
            fan_in = param.shape[1]
            std = math.sqrt(2.0 / fan_in) / math.sqrt(2 * n_layers)
            nn.init.normal_(param, std=std)

        else:
            # All other linear layers: He normal
            nn.init.kaiming_normal_(param, mode='fan_in')

    # RMSNorm weights: all ones (PyTorch default)
    for name, param in model.named_parameters():
        if 'norm' in name and param.dim() == 1:
            nn.init.ones_(param)

# Llama 2 7B: d_model=4096, n_layers=32
# Embedding std   = 1/sqrt(4096) = 0.0156
# Linear std      = sqrt(2/4096) = 0.0221
# Residual std    = 0.0221 / sqrt(64) = 0.00276
# RMSNorm weights = [1, 1, 1, ..., 1]
# Biases          = none (Llama has no biases)

Chapter 8: The Initialization Arena

Let's race them all.

You've learned five initialization strategies: Zero, Random, Xavier, He/Kaiming, and Orthogonal. Each has strengths and weaknesses — but reading about failure modes is one thing. Watching them happen in real time is another.

This simulation trains five identical networks on the same task, differing only in initialization. Drag the sliders to create the conditions where each method fails. You'll discover that He + ReLU isn't always king — context matters.

What to Discover

Experiment 1: Depth = 32, ReLU. Xavier's loss flatlines or oscillates wildly — its variance halves every layer through ReLU (the ½ factor). He stays stable because it compensates with 2/fan_in. This is the textbook case.

Experiment 2: Tanh activation. Switch to tanh. Now Xavier is the winner — it was designed for symmetric activations. He overshoots because its larger variance pushes activations into tanh's saturation regions.

Experiment 3: Depth = 64. At extreme depth, even He struggles. Orthogonal initialization keeps gradients flowing because it preserves norms exactly (no expansion, no shrinkage). This is why Orthogonal is popular in RNNs and very deep networks.

Experiment 4: Shallow (depth = 4), wide (fan_in = 1024). All methods converge to similar performance. Initialization matters less for shallow networks — the problem is self-correcting over a few layers. This is why you rarely hear about initialization for 3-layer MLPs.

Failure Mode Summary

Chapter 9: Cheat Sheet & Connections

You now understand the complete initialization toolkit — from the symmetry-breaking problem to the recipes used in every modern LLM. This chapter is your practical reference. No new concepts. Just the formulas, the decision guide, and the connections to where you go next.

Every Formula at a Glance

The Decision Flowchart

Follow the path that matches your situation:

Symbol Glossary

PyTorch One-Liner Reference

python
import torch.nn as nn

# Xavier for tanh/sigmoid
nn.init.xavier_normal_(layer.weight)
nn.init.xavier_uniform_(layer.weight)

# He/Kaiming for ReLU
nn.init.kaiming_normal_(layer.weight, mode='fan_in', nonlinearity='relu')
nn.init.kaiming_uniform_(layer.weight, mode='fan_in', nonlinearity='relu')

# Orthogonal for deep nets / RNNs
nn.init.orthogonal_(layer.weight, gain=1.0)

# Truncated normal for transformer embeddings
nn.init.trunc_normal_(layer.weight, std=0.02)

# LLM residual scaling (GPT-2 style)
nn.init.normal_(layer.weight, std=0.02 / math.sqrt(2 * n_layers))

Chapter Recap

Where to Go Next