Weight Initialization in Deep Neural Networks


Introduction

How Modern Deep Networks Start Learning?

Training deep neural networks is notoriously sensitive to how the weights are initialized. A poor initialization can cause gradients to vanish or explode, leading to divergence or painfully slow convergence. Conversely, a good initialization can speed up training, stabilize optimization, and improve generalization.

This article compares and contrasts classical and modern weight initialization methods, explains their mathematical motivation, and demonstrates their impact by plotting loss curves for the same network trained with different initializations.

We will cover:

  1. Why initialization matters?
  2. Key methods: Zero, Random Normal/Uniform, Xavier/Glorot, He/Kaiming, Orthogonal, and modern approaches
  3. Empirical comparison via experiments
  4. Practical guidelines for choosing initialization

Why Initialization Matters

Consider a fully connected neural network with input dimension $d_{in}$, output dimension $d_{out}$, and weights $W \in \mathbb{R}^{d_{out} \times d_{in}}$. During forward propagation, the pre-activation for layer $l$ is:

$$z^{(l)} = W^{(l)} x^{(l-1)} + b^{(l)}$$

If the weights are too large, activations and gradients will grow exponentially with depth (exploding gradients). If too small, signals shrink to near-zero (vanishing gradients). Both make training unstable.

Graph of 3D Loss Surface with Multiple Minima and Gradient Descent Paths
Title: 3D Loss Surface with Multiple Minima and Gradient Descent Paths
Source: AIML.com Research


Diagram of Initialization as a Starting Point on a Multi-Minima Loss Surface
Title: Initialization as a Starting Point on a Multi-Minima Loss Surface
Source: AIML.com Research

In the figures above, we see three initial points: Poor initialization (red), good initialization (green), and far initialization (blue). These initializations can come in different forms, but in this case, we see that the red point needs a much longer path and runs through saddle points, and depending on the direction it takes, may get stopped before a “good enough” local minima. The good initialization allowed your descent algorithm to take a direct path to a local minima. The far initialization started on a point that is too far from where we are interested to descend to, and may lead to inconclusive results.

Purpose of Initialization

A good initialization aims to maintain a roughly constant variance of activations and gradients across layers. This principle underlies the most popular initialization schemes.

Think of initialization as choosing the first point on a rugged loss surface. Optimization is then the process of descending from this point toward a local or global minimum. A poor initialization can start you far from a good basin or in a flat plateau where gradients vanish.  A well-chosen initialization starts you closer to a productive region, where gradient descent quickly converges.

$$z^{(l)} = W^{(l)} x^{(l-1)} + b^{(l)}$$

Below, different initializations are shown as starting points on a multi-minima loss landscape. 

Points near the steep slopes (good initializations) converge faster, while those in flat regions (bad initializations) struggle to descend.

Classical Initialization Methods

Random Normal Initialization

Initialize weights from a Gaussian distribution:

$$W_{ij} \sim \mathcal{N}(0, \sigma^2)$$

where $\sigma$ is a chosen standard deviation.

  • If $\sigma$ is too large → exploding activations
  • If too small → vanishing activations

This was common in early networks (Rumelhart et al., 1986). Issue: Sensitive to $\sigma$ choice and does not adapt to layer size.

Random Uniform Initialization

Weights are drawn from:

$$W_{ij} \sim \mathcal{U}(-a, a)$$

where $a$ is typically small (e.g., 0.01).

Similar issues as normal initialization: scale is arbitrary and does not account for input/output dimensions.

Modern Initialization Methods

Xavier / Glorot Initialization (2010)

Glorot & Bengio (2010) proposed balancing the variance of activations and backpropagated gradients. For a layer with $d_{in}$ inputs and $d_{out}$ outputs:

$$Var(W_{ij}) = \frac{2}{d_{in} + d_{out}}$$

Common implementations:

  • Uniform:
$$W_{ij} \sim \mathcal{U}\left(-\sqrt{\frac{6}{d_{in}+d_{out}}}, ; \sqrt{\frac{6}{d_{in}+d_{out}}}\right)$$


  • Normal:
$$ W_{ij} \sim \mathcal{N}\left(0, \frac{2}{d_{in}+d_{out}}\right) $$

Balances forward and backward variance, working best with tanh or sigmoid activation functions.  

Issue: Underestimates variance for ReLU activations (half inputs are zeroed).

He / Kaiming Initialization (2015)

Now standard for deep ReLU networks (ResNet, etc.), He et al. (2015) argued that only half of inputs survive the ReLU (positive values). To compensate, they scaled variance:

$$ Var(W_{ij}) = \frac{2}{d_{in}} $$

Implementations:

  • Uniform:
$$ W_{ij} \sim \mathcal{U}\left(-\sqrt{\frac{6}{d_{in}}}, ; \sqrt{\frac{6}{d_{in}}}\right) $$


  • Normal:
$$ W_{ij} \sim \mathcal{N}\left(0, \frac{2}{d_{in}}\right) $$


This scaling preserves variance across layers and prevents gradient shrinkage in very deep models.However, this initialization was built for ReLU-based networks to solve the vanishing gradient problem. It therefore assumes the input is centered around its mean. It can also lead to the “dying ReLU” problem, where ReLU neurons can get stuck and always output zero, effectively becoming permanently inactive.

Orthogonal Initialization

Proposed by Saxe et al. (2014). Initialize $W$ as an orthogonal matrix:

$$ W^\top W = I $$


Preserves norm of vectors across layers, avoiding vanishing/exploding gradients. Useful in Recurrent Neural Networks (RNNs).However, this technique can lead to computationally expensive updates and does not work well for all networks.

Modern Variants

  • LeCun Initialization: Designed for SELU activations (Klambauer et al., 2017). Variance:
$$ Var(W_{ij}) = \frac{1}{d_{in}} $$


  • Scaled Orthogonal: Combines orthogonal with variance scaling

Food for Thought: Why Not Initialize Everything to Zero?

Set all weights to zero:

$$ W_{ij} = 0 $$

Early networks (Rumelhart et al., 1986) experimented with zero weights, but this leads to the symmetry problem, where all neurons in a layer receive identical gradients, and therefore update identically. The model effectively learns one feature per layer, collapsing to a linear transformation. 


Loss curve behavior: No symmetry breaking → no diversity in neuron behavior → model stagnates with no real learning 

Empirical Demonstration

Let’s compare initializations on the same dataset (MNIST, simple MLP).

Experiment Setup

  • Dataset: MNIST (handwritten digits)
  • Model: 3-layer fully connected neural network with ReLU
  • Optimizer: SGD (lr=0.01)

The code snippet below can act as a template, regardless of your input dataset but as a general guide in initializing weights for your needs:

Python
import torch
iaiming for ReLU-like activations
        nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
        nn.init.zeros_(m.bias)

# Alternative examples:
#   nn.init.xavier_uniform_(m.weight)  # Good for tanh/sigmoid
#   nn.init.orthogonal_(m.weight)      # Often useful for RNNs / very deep nets

# 3) APPLY INITIALIZATION HERE (before optimizer/training)
model.apply(init_weights)

# 4) Then continue with optimizer/loss/loop as usual
optimizer = torch.optim.SGD(model.parameters(), lr=1e-2, momentum=0.9)
criterion = nn.CrossEntropyLoss()

# Training loop (sketch)
# for x, y in dataloader:
#     logits = model(x)
#     loss = criterion(logits, y)
#     optimizer.zero_grad()
#     loss.backward()
#     optimizer.step()
import torch.nn as nn
import torch.nn.functional as F

# 1) Define your model
class MLP(nn.Module):
    def __init__(self, d_in=784, d_hidden=256, d_out=10):
        super().__init__()
        self.fc1 = nn.Linear(d_in, d_hidden)
        self.fc2 = nn.Linear(d_hidden, d_hidden)
        self.fc3 = nn.Linear(d_hidden, d_out)
        self.act = nn.ReLU()  # Change to Tanh() to demo Xavier

    def forward(self, x):
        x = self.act(self.fc1(x))
        x = self.act(self.fc2(x))
        return self.fc3(x)

model = MLP()

# 2) Choose an initializer based on activation
def init_weights(m):
    if isinstance(m, nn.Linear):
        # He/K

The code used to run the experiment can be seen in the code snippet below:

Python
import osimport math
import time
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset, random_split
import matplotlib.pyplot as plt
import pandas as pd

# -------------------------
# Device (MPS if available)
# -------------------------
def get_device():
  if torch.backends.mps.is_available():
      return torch.device("mps")
  if torch.cuda.is_available():
      return torch.device("cuda")
  return torch.device("cpu")

device = get_device()
print("Using device:", device)

# -------------------------
# Data: MNIST or synthetic
# -------------------------
def load_mnist_or_synth(train_frac=0.9, seed=7):
  torch.manual_seed(seed)
  try:
      from torchvision import datasets, transforms
      transform = transforms.Compose([transforms.ToTensor()])
      ds = datasets.MNIST(root="./data", train=True, download=True, transform=transform)
      X = ds.data.float().view(-1, 28*28) / 255.0
      y = ds.targets.long()
      # normalize per-feature (optional but stabilizes synthetic fallback match)
      X = (X - X.mean(0)) / (X.std(0) + 1e-8)
      full = TensorDataset(X, y)
      n_train = int(train_frac * len(full))
      n_val = len(full) - n_train
      train_ds, val_ds = random_split(full, [n_train, n_val], generator=torch.Generator().manual_seed(seed))
      return train_ds, val_ds, 28*28, 10, "MNIST"
  except Exception as e:
      print("Falling back to synthetic data (MNIST-like). Reason:", str(e))
      # Synthetic: 10 classes, 28x28 flattened inputs
      n_per_class = 800
      n_classes = 10
      d = 28*28
      rng = np.random.default_rng(seed)
      centers = rng.normal(size=(n_classes, d))
      centers = centers / (np.linalg.norm(centers, axis=1, keepdims=True) + 1e-12)
      centers *= 3.0
      Xs, ys = [], []
      for c in range(n_classes):
          noise = rng.normal(0, 1.0, size=(n_per_class, d))
          Xs.append(centers[c] + noise)
          ys.append(np.full(n_per_class, c, int))
      X = np.vstack(Xs).astype(np.float32)
      y = np.concatenate(ys).astype(np.int64)
      # standardize
      X = (X - X.mean(0)) / (X.std(0) + 1e-8)
      # split
      n_total = X.shape[0]
      n_train = int(train_frac * n_total)
      Xtr, ytr = X[:n_train], y[:n_train]
      Xva, yva = X[n_train:], y[n_train:]
      train_ds = TensorDataset(torch.from_numpy(Xtr), torch.from_numpy(ytr))
      val_ds = TensorDataset(torch.from_numpy(Xva), torch.from_numpy(yva))
      return train_ds, val_ds, d, n_classes, "Synthetic"

train_ds, val_ds, D_IN, N_CLASS, tag = load_mnist_or_synth()

# -------------------------
# Dataloaders
# -------------------------
BATCH_SIZE = 256
train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True)
val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False)

# -------------------------
# Model & Initializations
# -------------------------
class MLP(nn.Module):
  def __init__(self, d_in, hidden1=256, hidden2=128, d_out=10, init_type="xavier"):
      super().__init__()
      self.fc1 = nn.Linear(d_in, hidden1)
      self.fc2 = nn.Linear(hidden1, hidden2)
      self.fc3 = nn.Linear(hidden2, d_out)
      self.relu = nn.ReLU(inplace=True)
      self.apply(lambda m: self._init_weights(m, init_type))

  @staticmethod
  def _init_weights(m, init_type):
      if isinstance(m, nn.Linear):
          if init_type == "zero":
              nn.init.constant_(m.weight, 0.0)
          elif init_type == "normal":
              nn.init.normal_(m.weight, mean=0.0, std=0.01)
          elif init_type == "xavier":
              nn.init.xavier_uniform_(m.weight)
          elif init_type == "he":
              nn.init.kaiming_uniform_(m.weight, nonlinearity="relu")
          else:
              nn.init.xavier_uniform_(m.weight)
          nn.init.constant_(m.bias, 0.0)

  def forward(self, x):
      x = self.relu(self.fc1(x))
      x = self.relu(self.fc2(x))
      return self.fc3(x)

def train_one(init_type="xavier", epochs=8, lr=0.05, weight_decay=1e-4):
  model = MLP(D_IN, d_out=N_CLASS, init_type=init_type).to(device)
  criterion = nn.CrossEntropyLoss()
  optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=weight_decay)

  train_losses, val_losses = [], []

  for ep in range(epochs):
      model.train()
      running = 0.0
      n_batches = 0
      for xb, yb in train_loader:
          xb = xb.to(device); yb = yb.to(device)
          optimizer.zero_grad(set_to_none=True)
          logits = model(xb)
          loss = criterion(logits, yb)
          loss.backward()
          optimizer.step()
          running += loss.item()
          n_batches += 1
      train_loss = running / max(1, n_batches)
      train_losses.append(train_loss)

      model.eval()
      vl_running = 0.0
      n_vb = 0
      correct = 0
      total = 0
      with torch.no_grad():
          for xb, yb in val_loader:
              xb = xb.to(device); yb = yb.to(device)
              logits = model(xb)
              loss = criterion(logits, yb)
              vl_running += loss.item()
              n_vb += 1
              pred = logits.argmax(dim=1)
              correct += (pred == yb).sum().item()
              total += yb.numel()
      val_loss = vl_running / max(1, n_vb)
      val_acc = correct / max(1, total)
      val_losses.append(val_loss)
      print(f"[{init_type}] Epoch {ep+1:02d}  train_loss={train_loss:.4f}  val_loss={val_loss:.4f}  val_acc={val_acc:.4f}")
  return {"train_loss": train_losses, "val_loss": val_losses}

# -------------------------
# Run experiments
# -------------------------
EPOCHS = 5  # Increase if you want more smooth curves on your machine
inits = ["zero", "normal", "xavier", "he"]
histories = {}
start = time.time()
for init in inits:
  histories[init] = train_one(init, epochs=EPOCHS, lr=0.08 if tag=="Synthetic" else 0.05)
end = time.time()
print(f"Total runtime: {end - start:.2f}s on {tag} data")

# -------------------------
# Save metrics table
# -------------------------
rows = []
for name, h in histories.items():
  for e, (tr, va) in enumerate(zip(h["train_loss"], h["val_loss"]), start=1):
      rows.append({"init": name, "epoch": e, "train_loss": tr, "val_loss": va})
df = pd.DataFrame(rows)
csv_path = "mps_init_metrics.csv"
df.to_csv(csv_path, index=False)
# -------------------------
# Plot aesthetics (article-ready)
# -------------------------
plt.rcParams.update({
  "figure.figsize": (8, 6),
  "axes.grid": True,
  "grid.alpha": 0.4,
  "axes.labelsize": 13,
  "axes.titlesize": 14,
  "xtick.labelsize": 12,
  "ytick.labelsize": 12,
  "legend.fontsize": 11,
  "lines.linewidth": 2.5,
  "lines.markersize": 6,
})

import seaborn as sns
sns.set_theme()

epochs = list(range(1, EPOCHS+1))

# Training loss
plt.figure()
for name in inits:
  plt.plot(epochs, histories[name]["train_loss"], marker="o", label=name)
plt.xlabel("Epoch")
plt.ylabel("Training Loss")
plt.title("Training Loss vs. Epoch for Different Initializations")
plt.legend(frameon=True, fancybox=True, shadow=False, ncol=2)
plt.tight_layout()
loss_train_path = "mps_init_training_loss.png"
plt.savefig(loss_train_path, dpi=200, bbox_inches="tight")
plt.show()

# Validation loss
plt.figure()
for name in inits:
  plt.plot(epochs, histories[name]["val_loss"], marker="o", label=name)
plt.xlabel("Epoch")
plt.ylabel("Validation Loss")
plt.title("Validation Loss vs. Epoch for Different Initializations")
plt.legend(frameon=True, fancybox=True, shadow=False, ncol=2)
plt.tight_layout()
loss_val_path = "mps_init_validation_loss.png"
plt.savefig(loss_val_path, dpi=200, bbox_inches="tight")
plt.show()

  • Initializations tested: Zero, Normal, Xavier, He

Diagram of Training Loss vs. Epoch for Different Initializations
Title: Training Loss vs. Epoch for Different Initializations
Source: AIML.com Research

Diagram of Validation Loss vs. Epoch for Different Initializations
Title: Validation Loss vs. Epoch for Different Initializations
Source: AIML.com Research

Figure: The He curve (red) closely tracks but slightly outperforms the Xavier curve (green). Both achieve stable convergence, but He shows faster/more stable early-stage descent. Zero initialization (blue) has the largest training and validation loss. Normal/Orthogonal initialization (orange) starts a little worse than He and Xavier but eventually converges with the rest. Of note, this is specific to the example code I ran. One may experience different results for different neural networks, data, environments, etc. For reproducibility, kindly refer to the code snippet below.

Observations

  • Zero Initialization:
    • As expected, the network with zero-initialized weights fails to learn. Both training and validation loss remain essentially flat at a high level (~2.3, i.e. random guessing). This reflects the symmetry problem that without distinct weight values, neurons are indistinguishable and gradients cannot break the tie

  • Normal Initialization (σ = 0.01):
    • The network does learn, but convergence starts slower compared to better schemes. Training loss decreases steadily, but the trajectory indicates sensitivity to the small variance choice, slightly larger or smaller σ could lead to vanishing or unstable gradients

  • Xavier Initialization:
    • Shows stable convergence with low training and validation loss, confirming its design for maintaining variance in both activations and backpropagated gradients. Though originally motivated for sigmoid/tanh, Xavier still performs competitively with ReLU

  • He Initialization:
    • Achieves the fastest and most consistent convergence among the tested schemes, with training and validation loss both quickly dropping close to zero. This aligns with the theoretical motivation: scaling weights by $\sqrt{2/d_{\text{in}}}$ compensates for ReLU’s property of zeroing out negative inputs

    • NOTE: While the He and Xavier loss curves nearly overlap, closer inspection reveals He converges slightly faster during early epochs and exhibits less oscillation. This behavior arises because He scaling better compensates for ReLU’s half-active neurons, maintaining stable variance throughout training

Key Takeaway

The empirical curves clearly highlight the role of initialization:

  • Bad initialization (zero) halts training entirely
  • Arbitrary random initialization (normal with fixed σ) can work, but is fragile
  • Variance-scaled schemes (Xavier, He) dramatically improve learning stability and speed
  • For modern ReLU networks, He initialization is the most effective choice

Practical Guidelines

  • Use He initialization for ReLU and variants (Leaky ReLU, ELU)
  • Use Xavier for sigmoid/tanh networks
  • Use LeCun for SELU (self-normalizing networks)
  • Use Orthogonal for RNNs or very deep architectures
  • Avoid Zero initialization
  • Stay up to date with research and what is considered standard because this industry moves extremely fast and some of these concepts are subject to change if proven otherwise

Conclusion

Weight initialization is not just a technical detail, it fundamentally shapes training dynamics. Modern deep learning relies heavily on variance-preserving schemes (Xavier, He) that maintain stable signal propagation.

The empirical results confirm the theory that bad initializations stall training, while carefully designed schemes accelerate convergence.

A rule of thumb: match your initialization to your activation function.

Videos

  • This video, “Weight Initialization for Deep Feedforward Neural Networks” by Assembly AI is a compact walkthrough of why initialization matters, with clean visuals on variance preservation. Good for getting the big picture before diving into formulas (Runtime: 4 mins)
YouTube video
Weight Initialization for Deep Feedforward Neural Networks by Assembly AI on Youtube
  • This video titled, “Various Weight Initialization Techniques in Neural Networks” by Krish Naik, is a practical overview of popular schemes and when to use each. Helpful if you want a quick “how to choose” summary
YouTube video
Various Weight Initialization Techniques in Neural Network by Krish Naik on Youtube
  • This video titled “Create Artificial Neural Networks using Weight Initialization Tricks” by Krish Naik, is a hands-on coding demo that implements different initializers and compares training behavior
YouTube video
Create Artificial Neural Network using Weight Initialization Tricks by Krish Naik on Youtube

References

  1. Glorot, X., & Bengio, Y. (2010). Understanding the difficulty of training deep feedforward neural networks. Proceedings of AISTATS. Link
  2. He, K., Zhang, X., Ren, S., & Sun, J. (2015). Delving deep into rectifiers: Surpassing human-level performance on ImageNet classification. ICCV. Link
  3. Saxe, A. M., McClelland, J. L., & Ganguli, S. (2014). Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. ICLR. Link
  4. Rumelhart, D. E., Hinton, G. E., & Williams, R. J. (1986). Learning representations by back-propagating errors. Nature. Link
  5. Klambauer, G., Unterthiner, T., Mayr, A., & Hochreiter, S. (2017). Self-normalizing neural networks. NeurIPS. Link
  1. What is Xavier Initialization?
  2. What is Kaiming Initialization?

Author: Roy Gabriel, PhD, Georgia Tech

Author

Help us improve this post by suggesting in comments below:

– modifications to the text, and infographics
– video resources that offer clear explanations for this question
– code snippets and case studies relevant to this concept
– online blogs, and research publications that are a “must read” on this topic

Leave the first comment

Partner Ad
Find out all the ways that you can
Contribute