Table of Contents
Why initialization matters
Weight initialization sets the starting point of learning. Before any data is seen, your network already defines a chain of matrix multiplications and nonlinearities. If the initial weights are too large, activations and gradients can explode as they move through layers. If they are too small, signals can shrink toward zero and learning becomes slow or stalls. Good initialization aims to keep the scale of activations and gradients roughly stable across depth so that early training is well behaved.
In a feedforward layer of the form $y = Wx + b$, the statistics of $y$ depend on both $x$ and $W$. If many layers are stacked, small mismatches in scale compound. The practical goal is not to pick a perfect set of weights, but to pick a distribution whose variance is appropriate for the layer width and the activation function used after it.
Intuition, variance propagation, and the core idea
A useful mental model is variance propagation. Suppose inputs to a layer have roughly zero mean and variance $\mathrm{Var}(x)$. With independent weights with zero mean and variance $\mathrm{Var}(W)$, the pre activation variance tends to scale like
$$
\mathrm{Var}(y) \approx \text{fan\_in} \cdot \mathrm{Var}(W) \cdot \mathrm{Var}(x),
$$
where $\text{fan\_in}$ is the number of input connections per neuron. If you want $\mathrm{Var}(y)$ to be similar to $\mathrm{Var}(x)$, you choose $\mathrm{Var}(W)$ proportional to $1/\text{fan\_in}$. Once you add nonlinearities such as ReLU, the variance changes again, so the best constant depends on the activation.
This is the origin of popular initializations that scale by fan in, fan out, or both.
Fan in and fan out in PyTorch terms
Fan in is the number of inputs to a unit, and fan out is the number of outputs.
For nn.Linear(in_features, out_features), $\text{fan\_in} = \text{in\_features}$ and $\text{fan\_out} = \text{out\_features}$.
For nn.Conv2d(in_channels, out_channels, kernel_size=kH x kW), each output channel sees in_channels kH kW inputs per spatial location, so $\text{fan\_in} = \text{in\_channels} \cdot k_H \cdot k_W$ and $\text{fan\_out} = \text{out\_channels} \cdot k_H \cdot k_W$.
PyTorch initialization utilities compute these values internally when you use functions in torch.nn.init.
Xavier, He, and common choices
Xavier, also called Glorot initialization, is designed for tanh like nonlinearities and tries to balance forward and backward signal. A common form sets the weight variance to
$$
\mathrm{Var}(W) = \frac{2}{\text{fan\_in} + \text{fan\_out}}.
$$
He, also called Kaiming initialization, is designed for ReLU family activations that zero out about half the inputs. A common form sets
$$
\mathrm{Var}(W) = \frac{2}{\text{fan\_in}}.
$$
In practice, the rule of thumb is to use Xavier for tanh or sigmoid, and He for ReLU or LeakyReLU. GELU often works well with He style variance scaling too, although many modern architectures rely on defaults that already make reasonable choices.
Do not initialize all weights to the same value, such as all zeros. This causes symmetry, meaning neurons in the same layer receive identical gradients and stay identical, so the network fails to learn useful diverse features.
Bias initialization
Biases are typically initialized to zero or a small constant. This is usually safe because biases do not create the same symmetry issue as weights.
There are a few special cases. For ReLU networks, a small positive bias can sometimes help avoid dead neurons early, but it is not a universal fix. For LSTMs and GRUs, some practitioners initialize forget gate biases to positive values, but the detailed rationale belongs in the sequence models chapter.
What PyTorch does by default
Many torch.nn modules come with reasonable default initializations. For example, nn.Linear uses a variant of Kaiming uniform initialization for weights and a uniform range for biases based on fan in. Convolution layers similarly use Kaiming based defaults. This means you can often train beginner level models without touching initialization.
Initialization becomes more important when you build custom modules, use unusual activations, create very deep networks, or observe unstable training where loss diverges or gradients vanish.
Custom initialization in PyTorch
You can override initialization by applying functions from torch.nn.init to parameters after constructing the model. The common pattern is to define a function that checks module types and initializes accordingly, then call model.apply(init_fn).
For example, Kaiming for linear and conv layers with ReLU can be expressed as Kaiming normal or uniform. If you use torch.nn.init.kaiming_normal_, you should specify the nonlinearity, such as "relu" or "leaky_relu", because the gain depends on it. Xavier has similar helpers like torch.nn.init.xavier_uniform_ and torch.nn.init.xavier_normal_.
A key detail is that these functions modify tensors in place, and you should run them with gradients disabled.
When manually initializing parameters, do it inside a torch.no_grad() block or by using in place init functions that are designed for initialization. You do not want initialization steps to be tracked by autograd.
Initializing embeddings and normalization layers
Embedding layers are often initialized from a normal distribution with a small standard deviation. PyTorch’s defaults are usually acceptable. For normalization layers such as BatchNorm and LayerNorm, the typical initialization is scale weights to ones and shift biases to zeros, which preserves identity behavior initially. Many modules already do this by default.
Diagnosing initialization related problems
If training loss becomes nan very early, or explodes within a few steps, initialization scale can be too large or combined with a too high learning rate. If loss barely changes and gradients are extremely small in early layers, initialization scale may be too small, or the network depth and activation choice may be causing vanishing gradients.
A practical sign is to inspect activation and gradient magnitudes for a single batch right after model creation. You are not looking for exact values, but you want to avoid values that are consistently near zero everywhere or extremely large.
If you change initialization and training becomes stable, but accuracy remains poor, do not keep tuning initialization as a substitute for data preprocessing, architecture choice, or learning rate selection. Initialization mainly helps optimization get started, it rarely fixes a mismatched model or pipeline.
Practical rules of thumb
Use PyTorch defaults for your first pass. If you explicitly choose, use He initialization for ReLU like activations and Xavier for tanh or sigmoid. Initialize biases to zero unless you have a specific reason. When you add new custom layers, ensure fan in and fan out are computed correctly, especially for convolutions. Keep in mind that initialization interacts with learning rate, so if you change one, you may need to revisit the other.