Offensive Security Researcher & Quantum Cryptography Analyst

Weight Initialization in Neural Networks: Concepts, Math, and Practice

Abdul Wahab Junaid
Explain the concept of weight initialization in neural networks

Introduction

It’s easy to overlook weight initialization as a mere technical formality — after all, the network is going to update its weights through training anyway, right? But early in my experimentation with deeper architectures, I learned the hard way that poor initialization can make a network essentially untrainable, no matter how good the architecture or optimizer is. Gradients either vanish into near-zero noise or explode into instability within the first few layers, and training never gets off the ground. This article explains why weight initialization matters so much, the mathematics behind the most widely used schemes, and practical guidance for choosing the right one.

Why Weight Initialization Matters

Neural networks are trained via backpropagation, which relies on computing gradients layer by layer using the chain rule. If weights are initialized poorly, the scale of activations and gradients can grow or shrink exponentially as they pass through many layers, leading to two classic problems:

Both problems are exacerbated in deep networks, where the compounding effect across many layers can turn even a mild imbalance into a severe one.

A Naive (Bad) Approach: All-Zero Initialization

If all weights are initialized to the same value (like zero), every neuron in a given layer computes the exact same output and receives the exact same gradient during backpropagation. This is called the symmetry problem — the network effectively behaves as if it had only one neuron per layer, regardless of how many neurons it actually has, since they all update identically.

$$ w_{ij}^{(l)} = 0 \quad \forall i,j,l \implies \text{all neurons in layer } l \text{ remain identical throughout training} $$

This is why weights must be initialized randomly — to break symmetry and allow different neurons to learn different features.

Another Naive Approach: Large Random Initialization

Simply using large random values (e.g., drawn from $\mathcal{N}(0, 1)$) without regard to layer size tends to cause activations to saturate quickly, particularly with sigmoid or tanh activation functions, pushing outputs toward the flat regions of these functions where gradients are near zero — a specific instance of the vanishing gradient problem.

The Variance-Preservation Principle

The key insight behind modern initialization schemes is that we want the variance of activations (and their gradients) to remain roughly constant as they pass through each layer of the network. If variance shrinks layer by layer, we get vanishing gradients; if it grows, we get exploding gradients.

Consider a single neuron computing $z = \sum_{i=1}^{n} w_i x_i$, where $n$ is the number of inputs. Assuming inputs $x_i$ and weights $w_i$ are independent with zero mean:

$$ \text{Var}(z) = \sum_{i=1}^{n} \text{Var}(w_i)\text{Var}(x_i) = n \cdot \text{Var}(w) \cdot \text{Var}(x) $$

For the variance of $z$ to equal the variance of $x$ (preserving scale across layers), we need:

$$ \text{Var}(w) = \frac{1}{n} $$

This simple derivation is the foundation for the Xavier/Glorot initialization scheme.

Xavier (Glorot) Initialization

Proposed by Xavier Glorot and Yoshua Bengio in 2010, this scheme accounts for both the forward pass (number of input units, $n_{in}$) and the backward pass (number of output units, $n_{out}$) to keep variance balanced in both directions:

Uniform variant:

$$ w \sim \mathcal{U}\left(-\sqrt{\frac{6}{n_{in}+n_{out}}}, \sqrt{\frac{6}{n_{in}+n_{out}}}\right) $$

Normal variant:

$$ w \sim \mathcal{N}\left(0, \frac{2}{n_{in}+n_{out}}\right) $$

Xavier initialization works well with symmetric activation functions like tanh and sigmoid, where the assumption of roughly linear behavior near zero holds reasonably well.

He (Kaiming) Initialization

ReLU activations, unlike tanh or sigmoid, zero out roughly half of their inputs (all negative values become zero). This effectively halves the variance passing through the layer, so Xavier initialization tends to underestimate the needed weight variance for ReLU-based networks. Kaiming He and colleagues (2015) proposed a correction:

$$ w \sim \mathcal{N}\left(0, \frac{2}{n_{in}}\right) $$

This accounts for the fact that only about half of ReLU’s outputs are nonzero on average, doubling the required variance compared to the Xavier scheme’s forward-pass-only version.

import torch.nn as nn

# PyTorch's default initialization for Linear layers uses a variant of Kaiming initialization
layer = nn.Linear(256, 128)
nn.init.kaiming_normal_(layer.weight, mode='fan_in', nonlinearity='relu')
nn.init.zeros_(layer.bias)

Diagram: Choosing an Initialization Scheme

flowchart TD
    A[Choose Weight Initialization] --> B{What Activation Function?}
    B -->|Sigmoid or Tanh| C[Xavier/Glorot Initialization]
    B -->|ReLU or Variants like Leaky ReLU| D[He/Kaiming Initialization]
    B -->|Self-Normalizing - SELU| E[LeCun Initialization]
    C --> F[Balanced Variance Across Layers]
    D --> F
    E --> F
    F --> G[Stable Gradients During Training]

LeCun Initialization

Used primarily with SELU (Scaled Exponential Linear Unit) activations, which are designed to be self-normalizing:

$$ w \sim \mathcal{N}\left(0, \frac{1}{n_{in}}\right) $$

Orthogonal Initialization

For recurrent neural networks (RNNs), where the same weight matrix is applied repeatedly across many time steps, orthogonal initialization is often used. Orthogonal matrices preserve the norm of vectors they’re multiplied with, which helps prevent the repeated matrix multiplications inherent in RNN unrolling from causing exponential growth or decay in activations and gradients.

rnn_layer = nn.RNN(input_size=128, hidden_size=256)
nn.init.orthogonal_(rnn_layer.weight_hh_l0)

Comparison Table of Initialization Schemes

SchemeFormulaBest ForKey Assumption
Zero initialization$w = 0$Never (breaks symmetry)N/A – fails immediately
Naive random (large variance)$w \sim \mathcal{N}(0,1)$Rarely used directlyNo variance control
Xavier/Glorot$\text{Var}(w) = \frac{2}{n_{in}+n_{out}}$Sigmoid, tanhRoughly linear activation near zero
He/Kaiming$\text{Var}(w) = \frac{2}{n_{in}}$ReLU, Leaky ReLUHalf of inputs zeroed by activation
LeCun$\text{Var}(w) = \frac{1}{n_{in}}$SELU (self-normalizing nets)Self-normalizing activation function
OrthogonalOrthogonal matrixRNNs, deep linear networksNorm preservation across repeated multiplication

Bias Initialization

Biases are typically initialized to zero, since the symmetry-breaking problem doesn’t apply to biases in the same way it does to weights (weights, not biases, determine whether neurons compute identical functions). One notable exception: in LSTM networks, the forget gate bias is often initialized to a small positive value (e.g., 1) to encourage the network to retain memory by default early in training, rather than aggressively forgetting.

Practical Code Example: Comparing Initializations

import torch
import torch.nn as nn

def build_model(init_scheme='he'):
    model = nn.Sequential(
        nn.Linear(784, 256),
        nn.ReLU(),
        nn.Linear(256, 128),
        nn.ReLU(),
        nn.Linear(128, 10)
    )

    for layer in model:
        if isinstance(layer, nn.Linear):
            if init_scheme == 'xavier':
                nn.init.xavier_normal_(layer.weight)
            elif init_scheme == 'he':
                nn.init.kaiming_normal_(layer.weight, nonlinearity='relu')
            elif init_scheme == 'zero':
                nn.init.zeros_(layer.weight)
            nn.init.zeros_(layer.bias)

    return model

# Compare activation statistics across schemes
for scheme in ['zero', 'xavier', 'he']:
    model = build_model(scheme)
    x = torch.randn(64, 784)
    with torch.no_grad():
        activations = model[0](x)
        print(f"{scheme}: activation std = {activations.std().item():.4f}")

Running this comparison typically shows that zero initialization produces a degenerate output (no useful signal), while Xavier and He initialization produce activations with a healthy, non-collapsing standard deviation.

Advantages of Proper Weight Initialization

Disadvantages and Limitations

Real-World Use Cases

Best Practices

Summary

Weight initialization determines the starting point of a neural network’s optimization process, and poor choices can make deep networks effectively untrainable due to vanishing or exploding gradients. The core principle behind modern schemes — Xavier/Glorot for sigmoid/tanh, He/Kaiming for ReLU-family activations, LeCun for SELU, and orthogonal initialization for recurrent networks — is preserving the variance of activations and gradients as they propagate through the network. While not a silver bullet on its own, proper initialization is a foundational technique that, combined with normalization and thoughtful architecture design, enables the training of the deep networks that power modern AI systems.

References

Exit mobile version