Weight Initialization › PyTorch Implementation
Language AI HandbookAll chapters Weight Initialization › PyTorch ImplementationOne section · Sections in this chapter12 / 19 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Modern deep learning frameworks provide built-in initialization functions. PyTorch offers the full suite of initialization methods through the torch.nn.init module. Understanding how to use these functions and when to apply them is an essential practical skill.
# Demonstrate PyTorch's built-in initialization functions
torch.manual_seed(42)
init_examples = {
"xavier_uniform_": lambda m: nn.init.xavier_uniform_(m.weight),
"xavier_normal_": lambda m: nn.init.xavier_normal_(m.weight),
"kaiming_uniform_": lambda m: nn.init.kaiming_uniform_(
m.weight, nonlinearity="relu"
),
"kaiming_normal_": lambda m: nn.init.kaiming_normal_(
m.weight, nonlinearity="relu"
),
"orthogonal_": lambda m: nn.init.orthogonal_(m.weight, gain=np.sqrt(2)),
}PyTorch built-in initialization statistics (256x512 layer, fan_in=512, fan_out=256): ----------------------------------------------------------------- xavier_uniform_ : std = 0.05111, range = [-0.0884, 0.0884] xavier_normal_ : std = 0.05110, range = [-0.2352, 0.2104] kaiming_uniform_ : std = 0.06252, range = [-0.1083, 0.1082] kaiming_normal_ : std = 0.06247, range = [-0.2800, 0.2481] orthogonal_ : std = 0.06250, range = [-0.2696, 0.2816]
The key PyTorch parameters for Kaiming (He) initialization are:
mode('fan_in'or'fan_out'): Controls whether to use the input dimension or output dimension for the variance calculation. Use'fan_in'(the default) to preserve forward pass variance, or'fan_out'to preserve backward pass variance.nonlinearity('relu','leaky_relu','tanh', etc.): Specifies the activation function so PyTorch computes the correct gain factor automatically.a(for Leaky ReLU): The negative slope parameter. PyTorch computes the gain as , matching our derived formula.
For Xavier initialization, the gain parameter multiplies the standard deviation. PyTorch provides nn.init.calculate_gain('tanh') and similar functions to look up the recommended gain for a given activation.
Applying Initialization to a Custom Network
In practice, you'll often want to apply initialization to an entire model at once. The cleanest pattern is to write an _init_weights method and call self.apply() to walk through all submodules.
class MLP(nn.Module):
"""A multi-layer perceptron with configurable initialization."""
def __init__(
self, input_size, hidden_sizes, output_size, activation="relu"
):
super().__init__()
all_sizes = [input_size] + hidden_sizes + [output_size]
self.layers = nn.ModuleList(
[
nn.Linear(all_sizes[i], all_sizes[i + 1])
for i in range(len(all_sizes) - 1)
]
)
self.activation = activation
self._init_weights()
def _init_weights(self):
"""Apply appropriate initialization based on activation function."""
for layer in self.layers:
if self.activation == "relu":
nn.init.kaiming_normal_(layer.weight, nonlinearity="relu")
elif self.activation == "tanh":
nn.init.xavier_normal_(
layer.weight, gain=nn.init.calculate_gain("tanh")
)
else:
nn.init.xavier_normal_(layer.weight)
nn.init.zeros_(layer.bias)
def forward(self, x):
for i, layer in enumerate(self.layers):
x = layer(x)
if i < len(self.layers) - 1:
if self.activation == "relu":
x = torch.relu(x)
elif self.activation == "tanh":
x = torch.tanh(x)
return xMLP initialization check: ReLU MLP, first layer std: 0.1424 Tanh MLP, first layer std: 0.1572 Expected ReLU std (He): 0.1414 Expected Tanh std (Xavier): 0.0937
The _init_weights pattern scales cleanly from simple MLPs to complex architectures. When building transformer models, you'd typically call self.apply(self._init_weights) to recursively initialize all submodules, then optionally scale specific weights (like GPT-2's projection weights) with a second pass.