Weight initialization

Deep dive

How it works

In: a tensor shape (fan_in, fan_out), the activation (gain) and an RNG · Out: initial weights with the chosen variance · State: none (RNG seed only) · Cost: O(P) once, zero per step

For y = W x with i.i.d. zero-mean weights and inputs, Var(y) = fan_in · Var(W) · E[x²].

1. Forward. ReLU zeroes half of a symmetric input, so E[relu(z)²] = Var(z)/2. Keeping Var(z) constant needs fan_in · Var(W) · ½ = 1, so Var(W) = 2/fan_in (He et al.).
2. Backward. The same argument on Wᵀ gives fan_out · Var(W) = 1. Glorot averages both: 2/(fan_in + fan_out).
3. Gain. Generalize as std = gain / √fan_mode. PyTorch calculate_gain: linear and sigmoid 1, ReLU √2, tanh 5/3, leaky ReLU √(2/(1+slope²)), SELU 3/4.
4. Fan for convolutions. fan_in = C_in · k_h · k_w, fan_out = C_out · k_h · k_w.
5. Residual depth. With pre-norm, x_{l+1} = x_l + F(norm(x_l)) adds an O(1) variance each block, so after L blocks Var(x) grows about L times (without normalization it compounds exponentially). Scaling each branch's last layer by 1/√L (GPT-2: 0.02/√(2·n_layer) for attention and MLP output projections), or initializing it to zero, keeps the stream at O(1).
6. Biases start at 0. Exceptions: LSTM forget-gate bias 1 (Jozefowicz et al.), and the output bias of an imbalanced classifier set to log(p/(1−p)) so the first steps don't just learn the prior.

depth 50 ReLU MLP, width 256, input rms 1:
  Var(W) = 1/fan_in  → final rms ≈ 2^-25 ≈ 3e-8   (vanishes)
  Var(W) = 2/fan_in  → final rms order of 1 (0.72 here: finite-width drift)
  Var(W) = 4/fan_in  → final rms ≈ 2^25 ≈ 3e7     (explodes)

Code

/// xorshift64* + Box-Muller: enough randomness for init, no crates.
struct Rng(u64);
impl Rng {
    fn uniform(&mut self) -> f64 {                     // [0, 1)
        self.0 ^= self.0 >> 12; self.0 ^= self.0 << 25; self.0 ^= self.0 >> 27;
        (self.0.wrapping_mul(0x2545F4914F6CDD1D) >> 11) as f64 / (1u64 << 53) as f64
    }
    fn normal(&mut self) -> f64 {
        let (u, v) = (1.0 - self.uniform(), self.uniform());
        (-2.0 * u.ln()).sqrt() * (std::f64::consts::TAU * v).cos()
    }
}

/// shape = [out, in, kh, kw...]: the receptive field multiplies both fans.
fn fans(shape: &[usize]) -> (f64, f64) {
    let rf: usize = shape[2..].iter().product();
    ((shape[1] * rf) as f64, (shape[0] * rf) as f64)
}
fn numel(shape: &[usize]) -> usize { shape.iter().product() }

fn glorot_uniform(r: &mut Rng, shape: &[usize], gain: f64) -> Vec<f32> {
    let (fi, fo) = fans(shape);
    let a = gain * (6.0 / (fi + fo)).sqrt();            // Var = a²/3 = 2/(fi+fo)
    (0..numel(shape)).map(|_| ((2.0 * r.uniform() - 1.0) * a) as f32).collect()
}

/// gain: ReLU √2, tanh 5/3, linear 1, SELU 3/4 (torch.nn.init.calculate_gain).
fn he_normal(r: &mut Rng, shape: &[usize], gain: f64) -> Vec<f32> {
    let std = gain / fans(shape).0.sqrt();
    (0..numel(shape)).map(|_| (r.normal() * std) as f32).collect()
}

/// N(0, std²) truncated at ±2σ, rescaled so the result really has standard deviation `std`.
fn trunc_normal(r: &mut Rng, n: usize, std: f64) -> Vec<f32> {
    let s = std / 0.87962566103423978;                  // std of N(0,1) restricted to [-2, 2]
    (0..n).map(|_| loop { let z = r.normal(); if z.abs() <= 2.0 { break (z * s) as f32; } }).collect()
}

/// Rows-orthonormal matrix (rows <= cols) by modified Gram-Schmidt on Gaussian rows.
fn orthogonal(r: &mut Rng, rows: usize, cols: usize) -> Vec<f64> {
    let mut q: Vec<f64> = (0..rows * cols).map(|_| r.normal()).collect();
    for i in 0..rows {
        for j in 0..i {
            let d: f64 = (0..cols).map(|c| q[i * cols + c] * q[j * cols + c]).sum();
            for c in 0..cols { q[i * cols + c] -= d * q[j * cols + c]; }
        }
        let n = (0..cols).map(|c| q[i * cols + c].powi(2)).sum::<f64>().sqrt();
        for c in 0..cols { q[i * cols + c] /= n; }
    }
    q
}

/// GPT-2: every weight N(0, 0.02²); attention-out and MLP-out get 0.02/√(2·n_layer).
fn gpt2_residual_std(std: f64, n_layer: usize) -> f64 { std / ((2 * n_layer) as f64).sqrt() }

/// y = x Wᵀ for x: [b, i], w: [o, i].
fn matmul_t(x: &[f32], w: &[f32], b: usize, i: usize, o: usize) -> Vec<f32> {
    let mut y = vec![0.0; b * o];
    for r in 0..b { for c in 0..o {
        y[r * o + c] = (0..i).map(|k| x[r * i + k] * w[c * i + k]).sum();
    } }
    y
}
fn rms(x: &[f32]) -> f64 { (x.iter().map(|&v| (v as f64).powi(2)).sum::<f64>() / x.len() as f64).sqrt() }

fn main() {
    let mut r = Rng(0x9E3779B97F4A7C15);
    let (b, w) = (64, 256);
    for var_scale in [1.0f64, 2.0, 4.0] {               // 50 ReLU layers, Var(W) = scale / fan_in
        let mut x: Vec<f32> = (0..b * w).map(|_| r.normal() as f32).collect();
        for _ in 0..50 {
            let wt = he_normal(&mut r, &[w, w], var_scale.sqrt());
            x = matmul_t(&x, &wt, b, w, w).into_iter().map(|v| v.max(0.0)).collect();
        }
        println!("Var(W) = {var_scale}/fan_in -> rms after 50 ReLU layers: {:.3e}", rms(&x));
    }
    let g = glorot_uniform(&mut r, &[256, 512], 1.0);
    let var = g.iter().map(|&v| (v as f64).powi(2)).sum::<f64>() / g.len() as f64;
    assert!((var / (2.0 / 768.0) - 1.0).abs() < 0.02);
    assert!((rms(&trunc_normal(&mut r, 1_000_000, 0.02)) - 0.02).abs() < 2e-4);
    let q = orthogonal(&mut r, 32, 64);
    let d01: f64 = (0..64).map(|c| q[c] * q[64 + c]).sum();
    assert!(d01.abs() < 1e-12);
    assert!((gpt2_residual_std(0.02, 12) - 0.02 / 24f64.sqrt()).abs() < 1e-15);
    // Pre-norm residual stream, 24 blocks of x += W2 relu(W1 rmsnorm(x)), He init in the branch.
    for scale in [1.0f64, 1.0 / 24f64.sqrt()] {
        let mut x: Vec<f32> = (0..b * w).map(|_| r.normal() as f32).collect();
        for _ in 0..24 {
            let mut h = x.clone();
            for row in h.chunks_mut(w) { let s = rms(row) as f32; row.iter_mut().for_each(|v| *v /= s); }
            let w1 = he_normal(&mut r, &[4 * w, w], 2f64.sqrt());
            let w2 = he_normal(&mut r, &[w, 4 * w], scale);
            let a: Vec<f32> = matmul_t(&h, &w1, b, w, 4 * w).into_iter().map(|v| v.max(0.0)).collect();
            let f = matmul_t(&a, &w2, b, 4 * w, w);
            x.iter_mut().zip(&f).for_each(|(xi, fi)| *xi += fi);
        }
        println!("branch scale {scale:.3}: residual rms after 24 blocks {:.2}", rms(&x));
    }
}

Complexity and constants

• Sampling P weights is O(P) once. For a 7B model, order of seconds on a GPU. Orthogonal init needs a QR per matrix, O(n³), which can take minutes for very wide layers on CPU.
• Deep-net budget: per-layer variance error ε compounds to (1+ε)^L. At L = 100 an error of 5% gives a factor of about 130.

Efficiency

• Cost: sampling P weights is O(P), once; a 7B model takes order of seconds on a GPU. Orthogonal init needs a QR per matrix, O(n³), minutes for very wide layers on CPU. The real cost is a bad init: diverged or stalled runs, and (1+ε)^L variance drift (about 130x at L = 100, ε = 5%).
• Levers:
  • initialize on device, or on the meta device and materialize shards — avoids the CPU RNG and host-to-device copy (order of) — models above a few billion parameters
  • each rank initializes only its own shard — memory per rank is P/world size — FSDP and tensor parallelism
  • variance-correct scales (Glorot, He, residual 1/√(2L)) — removes warm-up failures (He et al. 2015) — deep nets
  • μP parameterization — tune hyperparameters on a small proxy and transfer them (Yang et al. 2021) — avoids full-scale sweeps
  • truncated normal via inverse CDF — one RNG draw per weight, no rejection loop — large tensors
• Naive → standard → state of the art:
  • per-weight RNG calls in a CPU loop, then copy to GPU · slow, variance unchecked
  • vectorized per-tensor sampling with fan-in/out scaling · O(P), seconds for 7B
  • sharded meta-device init plus μP scaling so hyperparameters transfer · O(P/ranks), no full-scale tuning

Implementation notes

• Know the default you are getting. PyTorch nn.Linear and nn.Conv2d use kaiming_uniform_(a=√5), which is U(±1/√fan_in). Hugging Face transformers usually override it with N(0, initializer_range²), commonly 0.02.
• Truncated normal at ±2σ (TensorFlow truncated_normal, torch.nn.init.trunc_normal_) avoids rare huge weights. Note that trunc_normal_ takes absolute bounds a, b (default ±2), not multiples of std: with std 0.02 it does almost nothing.
• Embeddings: N(0, 1) (PyTorch nn.Embedding default) is far too large when tied to the output layer of a transformer. Most LLM codebases use std 0.02 or 1/√d.
• Zero-init options: the last BN γ in each ResNet block (Goyal et al.), the last projection of each residual branch, adapter or LoRA B matrices, and the new head of a fine-tune. Never zero-init all layers: identical neurons get identical gradients and never break symmetry.
• muP (Yang et al.) scales init and per-layer learning rates with width, so hyperparameters tuned on a small model transfer to a large one.

Pitfalls

• He init with tanh or sigmoid: tanh saturates and gradients vanish.
• Using fan_in for transposed convs or weights stored as (in, out): your fans() must match the layout, or Var is off by fan_out/fan_in.
• Changing init without retuning learning rate: the effective step size relative to weight norm changes.
• Very large output-layer init in a classifier: initial loss far above ln K. Check the first-step loss against ln K.
• Forgetting the 1/√N residual scaling in a homemade transformer: early loss spikes at depth 24+ that look like a learning-rate problem.

Variants and alternatives

In production

torch.nn.init (kaiming_normal_, xavier_uniform_, orthogonal_, trunc_normal_), Flax and Keras variance_scaling. GPT-2 and nanoGPT (0.02, residual projection scaled), Llama-style codebases (0.02 normal), timm ResNets (zero-init last BN γ).

Implementations

Sources

• Glorot and Bengio, Understanding the Difficulty of Training Deep Feedforward Neural Networks, AISTATS 2010. paper
• He, Zhang, Ren, Sun, Delving Deep into Rectifiers, ICCV 2015. paper
• LeCun, Bottou, Orr, Müller, Efficient BackProp, 1998.
• Saxe, McClelland, Ganguli, Exact Solutions to the Nonlinear Dynamics of Learning in Deep Linear Neural Networks, ICLR 2014. paper
• Radford et al., Language Models are Unsupervised Multitask Learners (GPT-2), 2019. paper
• Zhang, Dauphin, Ma, Fixup Initialization: Residual Learning Without Normalization, ICLR 2019. paper
• Goyal et al., Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour, 2017. paper
• Yang et al., Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer, 2021. paper
• Jozefowicz, Zaremba, Sutskever, An Empirical Exploration of Recurrent Network Architectures, ICML 2015. paper