Weight initialization: why your deep net trains or dies before step one

A deep network does not begin learning from a neutral state. Before the optimizer takes its first step, the initial weights have already determined:

• How large each layer’s activations are

• How much information survives through the network

• Whether gradients can reach early layers

• Whether floating-point values remain finite

• Whether identical neurons can learn different features

• How aggressively the first optimizer updates will move the model

A poor initialization can make a valid architecture appear broken. Loss may remain flat, gradients may become effectively zero, or activations may grow until they overflow. Conversely, a suitable initialization can turn the same model, dataset, optimizer, and learning rate into a healthy training run.

That is why weight initialization is not cosmetic randomness. It is the first stability decision in the training pipeline.

What initialization is trying to preserve

Consider a linear layer followed by an activation:

z = xWᵀ + b

h = activation(z)

Each output unit combines fan_in input values. When the inputs and weights are roughly independent and centered near zero, the variance of the output depends on three main quantities:

• The variance of the incoming activations

• The variance of the weights

• The number of incoming connections, or fan_in

If every incoming connection contributes independent noise, adding more connections increases the output variance. A layer with 1,024 inputs therefore cannot safely use the same weight scale as a layer with 8 inputs.

Initialization methods compensate by making individual weights smaller as fan_in grows.

The objective is not to make every activation numerically identical. The practical objective is to keep the scale within a useful range as signals move through many layers.

For the forward pass, you generally want:

• Activations that remain finite

• Standard deviations that do not collapse toward zero

• Standard deviations that do not grow exponentially

• Nonlinearities that operate in responsive regions rather than saturated regions

For the backward pass, you want the same broad property for gradients. A gradient should be able to travel from the loss to early layers without disappearing or exploding.

How a bad scale compounds with depth

Suppose each layer multiplies the activation standard deviation by approximately 0.7. One layer is harmless. Thirty layers are not:

0.7³⁰ ≈ 0.0000225

The useful signal has nearly disappeared.

Now suppose each layer multiplies the scale by 1.5:

1.5³⁰ ≈ 191,751

Values can become enormous before the model has processed a single training example.

Real neural networks do not behave as cleanly as those scalar examples. Activations, correlations, biases, normalization layers, residual connections, and changing widths all affect the result. The compounding principle still holds: a modest per-layer scale error becomes a serious network-level error.

This creates two classic failure modes.

Vanishing activations and gradients

Weights that are too small produce progressively smaller activations. Eventually:

• Hidden representations become nearly constant

• Early-layer gradients approach floating-point noise

• The optimizer updates only the final layers meaningfully

• Loss appears frozen even though backward() completes successfully

A vanishing-gradient run is especially deceptive because it may contain no exceptions or NaN values. The training loop runs normally; it simply does not learn.

Exploding activations and gradients

Weights that are too large produce rapidly growing activations. This may cause:

• Extreme logits

• Unstable losses

• Huge gradient norms

• Repeated gradient clipping

• inf or NaN values

• Overflow that appears only under mixed precision

Gradient clipping can limit the final gradient norm, but it cannot repair a forward pass whose internal representations are already badly scaled. Clipping is a guardrail, not a substitute for initialization.

Why initializing every weight to zero fails

Zero initialization sounds stable because it cannot explode. It creates a different problem: symmetry.

Imagine two neurons in the same layer. If they begin with identical weights, receive the same inputs, and participate in the same computation, they receive identical gradients. After the first update, they remain identical. The same remains true after every later update.

Instead of learning two distinct features, the layer behaves like duplicated copies of one feature.

Biases can often begin at zero because random weights already break the symmetry between neurons. Weight matrices generally need non-identical values.

The lesson is subtle:

• Weight scale should be controlled.

• Weight values should not all be identical.

• Randomness is useful because it breaks symmetry.

• The distribution and variance of that randomness determine signal stability.

Xavier initialization

Xavier initialization, also called Glorot initialization, balances the incoming and outgoing dimensions of a layer.

For a normal distribution, PyTorch uses:

std = gain × √(2 / (fan_in + fan_out))

For a uniform distribution, the values are sampled between negative and positive bounds:

bound = gain × √(6 / (fan_in + fan_out))

The use of both fan_in and fan_out is a compromise between preserving forward activations and preserving backward gradients.

Xavier initialization is a natural starting point for:

• Linear layers followed by no strong variance-changing nonlinearity

• Tanh networks, with the appropriate gain

• Sigmoid networks, although deep sigmoid networks still face saturation risks

• Output projections without a ReLU after them

• Architectures whose established recipe explicitly specifies Xavier initialization

PyTorch provides xavier_uniform_ and xavier_normal_, along with calculate_gain() for supported nonlinearities. Its current initialization functions execute in torch.no_grad() mode, so initialization operations are not recorded by autograd.

A direct example looks like this:

from torch import nn


layer = nn.Linear(256, 128)

nn.init.xavier_uniform_(
    layer.weight,
    gain=nn.init.calculate_gain("tanh"),
)
nn.init.zeros_(layer.bias)

The gain matters. “Use Xavier” is incomplete advice unless the following activation is also considered.

He initialization for ReLU networks

ReLU removes every negative pre-activation:

ReLU(x) = max(0, x)

With a roughly symmetric input distribution, approximately half the values become zero. That changes the signal’s second moment and calls for a larger starting variance than plain Xavier initialization with gain 1.

He initialization, also called Kaiming initialization, compensates for this behavior. For ReLU with fan_in mode, the normal-distribution scale is:

std = √(2 / fan_in)

The corresponding uniform distribution uses a bound equivalent to:

bound = √(6 / fan_in)

In PyTorch:

from torch import nn


layer = nn.Linear(256, 128)

nn.init.kaiming_normal_(
    layer.weight,
    mode="fan_in",
    nonlinearity="relu",
)
nn.init.zeros_(layer.bias)

Use fan_in when the main objective is preserving forward activation scale. fan_out is sometimes selected when preserving backward-gradient scale is more important, particularly in specialized convolutional designs.

For leaky ReLU, pass the actual negative slope:

from torch import nn


negative_slope = 0.1
layer = nn.Linear(256, 128)

nn.init.kaiming_uniform_(
    layer.weight,
    a=negative_slope,
    mode="fan_in",
    nonlinearity="leaky_relu",
)

Passing an incorrect slope changes the computed gain and therefore changes the initial variance.

A surprising equivalence: Xavier can become He

For a square layer where fan_in = fan_out, Xavier normal initialization with a ReLU gain becomes mathematically identical in scale to He normal initialization.

Xavier with ReLU gain:

std = √2 × √(2 / (fan_in + fan_out))

For a square layer:

std = √2 × √(2 / (2 × fan_in))

This simplifies to:

std = √(2 / fan_in)

That is the He normal scale.

This explains why the runnable experiment later in this article produces nearly identical activation statistics for xavier_normal_ with a ReLU gain and kaiming_normal_ across square hidden layers. The methods diverge when layer dimensions are rectangular or when different modes and gains are selected.

Current PyTorch defaults are not simply “He for ReLU”

As of PyTorch 2.13, released in July 2026, the stable initialization APIs include Xavier, Kaiming, orthogonal, sparse, Dirac, truncated normal, and standard distribution-filling operations.

The important practical detail is that built-in modules choose their own defaults.

nn.Linear

In PyTorch 2.13, nn.Linear initializes its weights from:

uniform(-1 / √fan_in, +1 / √fan_in)

The implementation calls:

nn.init.kaiming_uniform_(weight, a=math.sqrt(5))

That call may look like ordinary He initialization, but a=√5 is deliberately chosen so that the resulting bound becomes 1 / √fan_in. It is a historical module default, not the same configuration as kaiming_uniform_(..., nonlinearity="relu"). Bias values use the same ±1 / √fan_in bound.

This means the following model does not automatically receive explicit He-for-ReLU initialization merely because a ReLU follows each linear layer:

from torch import nn


model = nn.Sequential(
    nn.Linear(256, 256),
    nn.ReLU(),
    nn.Linear(256, 256),
    nn.ReLU(),
)

PyTorch’s default is often serviceable, especially in modest networks or architectures stabilized by normalization and residual paths. It should not be confused with an architecture-aware initialization policy.

Convolutional layers

The current Conv1d, Conv2d, and Conv3d implementation also uses Kaiming uniform with a=√5, followed by bias initialization based on 1 / √fan_in. PyTorch computes convolutional fan values from the channel, group, and kernel dimensions rather than treating the kernel as a simple two-dimensional matrix.

That is one reason to use the framework’s initialization helpers rather than manually calculating a scale from weight.shape[1].

nn.Transformer

The top-level nn.Transformer reset routine applies Xavier uniform initialization to parameters with more than one dimension. PyTorch’s source also warns that TransformerEncoder layers begin with the same parameter values because the supplied encoder layer is cloned; the documentation recommends manually initializing the layers after constructing the encoder stack.

That warning is easy to miss when relying only on familiar module names.

A practical initialization policy

A useful policy should follow the architecture rather than blindly replacing every parameter.

from __future__ import annotations

import torch
from torch import nn


def initialize_relu_network(module: nn.Module) -> None:
    if isinstance(module, nn.Linear):
        nn.init.kaiming_normal_(
            module.weight,
            mode="fan_in",
            nonlinearity="relu",
        )
        if module.bias is not None:
            nn.init.zeros_(module.bias)

    elif isinstance(module, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
        nn.init.kaiming_normal_(
            module.weight,
            mode="fan_in",
            nonlinearity="relu",
        )
        if module.bias is not None:
            nn.init.zeros_(module.bias)


model = nn.Sequential(
    nn.Linear(64, 256),
    nn.ReLU(),
    nn.Linear(256, 256),
    nn.ReLU(),
    nn.Linear(256, 1),
)

model.apply(initialize_relu_network)

nn.init.xavier_uniform_(model[-1].weight)
nn.init.zeros_(model[-1].bias)

Notice that the final output layer is handled separately. The hidden layers are followed by ReLU, so they receive He initialization. The output layer has no ReLU after it, so Xavier is a more defensible general-purpose choice.

This is not a universal law. Some models intentionally use:

• Small normal initialization for output heads

• Zero initialization for residual-branch output projections

• Orthogonal recurrent matrices

• Truncated normal distributions

• Architecture-specific depth scaling

• Initialization values inherited from a pretrained checkpoint

The correct rule is: match the initialization to the layer’s role, activation, shape, and surrounding architecture.

Activation-to-initialization guide

A practical starting guide is:

Following operation	Reasonable starting point
ReLU	Kaiming with nonlinearity="relu"
Leaky ReLU	Kaiming with the real negative slope
Tanh	Xavier with calculate_gain("tanh")
Sigmoid	Xavier with gain 1
Linear output	Xavier or the architecture’s documented recipe
SELU self-normalizing network	Kaiming with nonlinearity="linear" as documented by PyTorch
GELU or SiLU	Use the architecture’s established recipe; PyTorch has no dedicated calculate_gain() entry

PyTorch’s current gain table supports linear operations, sigmoid, tanh, ReLU, leaky ReLU, and SELU. Its SELU documentation specifically notes that self-normalizing networks should use the linear gain to obtain the intended 1 / fan_in variance.

For GELU and SiLU, avoid inventing a gain value and presenting it as a framework standard. Modern transformer and vision architectures often rely on a complete recipe involving normalization placement, residual scaling, width, depth, and sometimes truncated distributions. Preserve that recipe unless you are deliberately running an ablation.

Runnable lab: watch activations vanish and explode

Create an isolated environment and install PyTorch:

python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install torch

Save the following as weight_init_lab.py:

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Iterable

import torch
from torch import nn


torch.set_num_threads(1)


@dataclass(frozen=True)
class LayerStats:
    layer: int
    activation_std: float
    zero_fraction: float


@dataclass(frozen=True)
class TrainResult:
    initial_loss: float
    final_loss: float
    first_layer_grad_norm: float


class DeepMLP(nn.Module):
    def __init__(
        self,
        input_dim: int,
        width: int,
        depth: int,
        output_dim: int,
        *,
        bias: bool,
    ) -> None:
        super().__init__()
        self.hidden = nn.ModuleList()
        current_dim = input_dim

        for _ in range(depth):
            self.hidden.append(
                nn.Linear(current_dim, width, bias=bias)
            )
            current_dim = width

        self.output = nn.Linear(
            current_dim,
            output_dim,
            bias=bias,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for layer in self.hidden:
            x = torch.relu(layer(x))
        return self.output(x)


def initialize(model: nn.Module, scheme: str) -> None:
    for module in model.modules():
        if not isinstance(module, nn.Linear):
            continue

        if scheme == "tiny":
            nn.init.normal_(
                module.weight,
                mean=0.0,
                std=0.01,
            )
        elif scheme == "huge":
            nn.init.normal_(
                module.weight,
                mean=0.0,
                std=1.0,
            )
        elif scheme == "xavier_relu":
            nn.init.xavier_normal_(
                module.weight,
                gain=nn.init.calculate_gain("relu"),
            )
        elif scheme == "he":
            nn.init.kaiming_normal_(
                module.weight,
                mode="fan_in",
                nonlinearity="relu",
            )
        else:
            raise ValueError(
                f"Unknown initialization scheme: {scheme}"
            )

        if module.bias is not None:
            nn.init.zeros_(module.bias)


def finite_std(tensor: torch.Tensor) -> float:
    finite = tensor[torch.isfinite(tensor)]
    if finite.numel() == 0:
        return math.inf
    return finite.std(unbiased=False).item()


def collect_activation_stats(
    model: DeepMLP,
    x: torch.Tensor,
) -> list[LayerStats]:
    stats: list[LayerStats] = []

    with torch.no_grad():
        for index, layer in enumerate(model.hidden, start=1):
            x = torch.relu(layer(x))
            stats.append(
                LayerStats(
                    layer=index,
                    activation_std=finite_std(x),
                    zero_fraction=(
                        (x == 0).float().mean().item()
                    ),
                )
            )

    return stats


def print_selected_stats(
    scheme: str,
    stats: list[LayerStats],
) -> None:
    selected = {1, 5, 10, 20, 30}

    print(f"\n{scheme}")
    print("layer  activation_std  zero_fraction")

    for item in stats:
        if item.layer in selected:
            print(
                f"{item.layer:5d}  "
                f"{item.activation_std:14.3e}  "
                f"{item.zero_fraction:13.3f}"
            )


def run_activation_experiment() -> None:
    print("=== Forward-signal diagnostic ===")

    generator = torch.Generator().manual_seed(7)
    x = torch.randn(512, 128, generator=generator)

    for scheme in (
        "tiny",
        "huge",
        "xavier_relu",
        "he",
    ):
        torch.manual_seed(11)

        model = DeepMLP(
            input_dim=128,
            width=128,
            depth=30,
            output_dim=1,
            bias=False,
        )
        initialize(model, scheme)

        stats = collect_activation_stats(model, x)
        print_selected_stats(scheme, stats)


def make_regression_data() -> tuple[
    torch.Tensor,
    torch.Tensor,
]:
    generator = torch.Generator().manual_seed(2026)
    x = torch.randn(4096, 32, generator=generator)

    y = (
        1.4 * torch.sin(x[:, 0])
        + 0.7 * x[:, 1] * x[:, 2]
        - 0.8 * torch.tanh(x[:, 3] - x[:, 4])
        + 0.3 * x[:, 5].square()
    )

    y = (y - y.mean()) / y.std()
    return x, y


def train_once(
    scheme: str,
    steps: int = 500,
) -> TrainResult:
    torch.manual_seed(11)

    model = DeepMLP(
        input_dim=32,
        width=64,
        depth=12,
        output_dim=1,
        bias=False,
    )
    initialize(model, scheme)

    x, y = make_regression_data()

    optimizer = torch.optim.SGD(
        model.parameters(),
        lr=0.01,
        momentum=0.9,
    )
    loss_fn = nn.MSELoss()

    with torch.no_grad():
        initial_loss = loss_fn(
            model(x).squeeze(-1),
            y,
        ).item()

    first_layer_grad_norm = math.nan

    for step in range(steps):
        optimizer.zero_grad(set_to_none=True)

        prediction = model(x).squeeze(-1)
        loss = loss_fn(prediction, y)
        loss.backward()

        if step == 0:
            gradient = model.hidden[0].weight.grad
            if gradient is None:
                raise RuntimeError(
                    "First-layer gradient was not computed"
                )
            first_layer_grad_norm = (
                gradient.norm().item()
            )

        torch.nn.utils.clip_grad_norm_(
            model.parameters(),
            max_norm=10.0,
        )
        optimizer.step()

    with torch.no_grad():
        final_loss = loss_fn(
            model(x).squeeze(-1),
            y,
        ).item()

    return TrainResult(
        initial_loss=initial_loss,
        final_loss=final_loss,
        first_layer_grad_norm=first_layer_grad_norm,
    )


def run_training_experiment(
    schemes: Iterable[str],
) -> None:
    print(
        "\n=== Same run, only initialization changes ==="
    )

    for scheme in schemes:
        result = train_once(scheme)

        print(
            f"{scheme:12s} "
            f"initial={result.initial_loss:.6f} "
            f"final={result.final_loss:.6f} "
            f"first_grad_norm="
            f"{result.first_layer_grad_norm:.3e}"
        )


if __name__ == "__main__":
    print(f"PyTorch {torch.__version__}")
    run_activation_experiment()
    run_training_experiment(
        ("tiny", "xavier_relu", "he")
    )

Run it:

python weight_init_lab.py

One fixed-seed CPU run produced the following representative output. Exact digits can vary with PyTorch version, hardware, and floating-point implementation.

PyTorch 2.10.0+cpu
=== Forward-signal diagnostic ===

tiny
layer  activation_std  zero_fraction
    1       6.575e-02          0.499
    5       3.252e-06          0.414
   10       1.238e-11          0.485
   20       1.360e-22          0.490
   30       9.395e-34          0.478

huge
layer  activation_std  zero_fraction
    1       6.575e+00          0.499
    5       3.252e+04          0.414
   10       1.238e+09          0.485
   20       1.360e+18          0.490
   30       9.395e+26          0.478

xavier_relu
layer  activation_std  zero_fraction
    1       8.218e-01          0.499
    5       9.924e-01          0.414
   10       1.153e+00          0.485
   20       1.179e+00          0.490
   30       7.589e-01          0.478

he
layer  activation_std  zero_fraction
    1       8.218e-01          0.499
    5       9.924e-01          0.414
   10       1.153e+00          0.485
   20       1.179e+00          0.490
   30       7.589e-01          0.478

=== Same run, only initialization changes ===
tiny         initial=0.999756 final=0.999756 first_grad_norm=2.137e-16
xavier_relu  initial=1.867352 final=0.031988 first_grad_norm=5.444e-01
he           initial=1.662760 final=0.024988 first_grad_norm=3.531e-01

Reading the activation results

The tiny initialization begins with an activation standard deviation around 0.066. By layer 10, it has fallen to approximately 1.24 × 10⁻¹¹. By layer 30, the signal is effectively gone.

The network is not merely producing “small values.” Every early feature must influence the output through a chain of tiny multiplications. The corresponding gradients become too small to drive useful updates.

The huge initialization has the opposite behavior. Its standard deviation grows from roughly 6.6 to 9.4 × 10²⁶. A deeper network, a wider network, mixed precision, or a more aggressive loss could push those values into non-finite territory.

Xavier with ReLU gain and He initialization keep the activation standard deviation near order one. The values fluctuate because this is one random finite-width network rather than an infinite-width theoretical model, but they neither vanish nor explode.

The similar zero fractions are also instructive. ReLU sets roughly half the activations to zero in every run. Looking only at sparsity would miss the scale failure. Diagnostics need both distribution shape and magnitude.

Cherry on the cake: one training run rescued by changing only initialization

The training section is a controlled rescue story.

The following remained unchanged:

• Synthetic regression dataset

• Data seed

• Model architecture

• Model depth and width

• Absence of biases

• Loss function

• SGD optimizer

• Learning rate

• Momentum

• Number of optimization steps

• Gradient-clipping threshold

Only the initialization scheme changed.

With std=0.01, the first layer’s initial gradient norm was approximately:

2.137 × 10⁻¹⁶

After 500 steps, the loss was unchanged to six decimal places:

0.999756 → 0.999756

The training loop was valid. The gradients were computed. The optimizer stepped. The network was nevertheless dead for practical purposes.

Changing only the initialization to He produced a first-layer gradient norm around 0.353. The final loss fell to approximately:

0.024988

Xavier with ReLU gain also rescued the run, reaching approximately 0.031988.

This is why initialization bugs are frequently misdiagnosed as optimizer problems. A team may try larger learning rates, AdamW, more training steps, loss rescaling, or gradient clipping when the real issue existed before step one.

Diagnose initialization before a long run

A short diagnostic pass can expose most severe initialization failures.

At minimum, inspect:

• Activation mean

• Activation standard deviation

• Minimum and maximum

• Fraction of exact zeros

• Fraction of finite values

• Gradient norm for each parameter

• Weight norm

• Update-to-weight ratio after the first optimizer step

Forward hooks make this convenient:

from __future__ import annotations

import torch
from torch import nn


def attach_activation_hooks(
    model: nn.Module,
) -> list[torch.utils.hooks.RemovableHandle]:
    handles: list[
        torch.utils.hooks.RemovableHandle
    ] = []

    def report(
        name: str,
    ):
        def hook(
            module: nn.Module,
            inputs: tuple[torch.Tensor, ...],
            output: torch.Tensor,
        ) -> None:
            if not isinstance(output, torch.Tensor):
                return

            detached = output.detach()
            finite = torch.isfinite(detached)

            finite_fraction = (
                finite.float().mean().item()
            )
            zero_fraction = (
                (detached == 0).float().mean().item()
            )

            if finite.any():
                values = detached[finite]
                mean = values.mean().item()
                std = values.std(
                    unbiased=False
                ).item()
            else:
                mean = float("nan")
                std = float("nan")

            print(
                f"{name:30s} "
                f"mean={mean: .3e} "
                f"std={std: .3e} "
                f"zeros={zero_fraction:.3f} "
                f"finite={finite_fraction:.3f}"
            )

        return hook

    for name, module in model.named_modules():
        if isinstance(
            module,
            (
                nn.Linear,
                nn.Conv1d,
                nn.Conv2d,
                nn.Conv3d,
            ),
        ):
            handles.append(
                module.register_forward_hook(
                    report(name)
                )
            )

    return handles

Use it for one representative batch, then remove the hooks:

handles = attach_activation_hooks(model)

with torch.no_grad():
    _ = model(example_batch)

for handle in handles:
    handle.remove()

Do not leave verbose diagnostic hooks attached during full training. They introduce synchronization, output, and bookkeeping overhead.

Measure gradients by layer

A model can have acceptable forward activations while its backward signal remains unhealthy. After one loss computation, inspect gradient norms:

from __future__ import annotations

from torch import nn


def report_gradient_norms(
    model: nn.Module,
) -> None:
    for name, parameter in model.named_parameters():
        if parameter.grad is None:
            print(f"{name:40s} grad=None")
            continue

        gradient_norm = parameter.grad.norm().item()
        weight_norm = parameter.detach().norm().item()

        ratio = gradient_norm / max(
            weight_norm,
            1e-12,
        )

        print(
            f"{name:40s} "
            f"weight={weight_norm:.3e} "
            f"grad={gradient_norm:.3e} "
            f"grad/weight={ratio:.3e}"
        )

The exact healthy range depends on the architecture and optimizer. The pattern matters more than a universal threshold.

Warning signs include:

• Early-layer gradients many orders of magnitude below later-layer gradients

• Non-finite gradients

• Nearly identical gradients across supposedly distinct units

• Gradient norms that grow exponentially toward the input

• Every gradient being exactly zero

• Large gradients paired with activations that already contain inf

Common initialization mistakes

Applying custom initialization after loading a checkpoint

This silently destroys pretrained weights:

model.load_state_dict(checkpoint)
model.apply(initialize_relu_network)

The correct order for training from a checkpoint is usually to construct the model, load the checkpoint, and initialize only newly added modules.

Reinitializing tied weights twice

model.apply() visits modules, not conceptual parameters. In architectures with tied or shared parameters, two modules may reference the same underlying tensor. A generic initialization callback can overwrite that tensor more than once.

Check parameter identities when weight tying is intentional.

Using .data for initialization

Avoid patterns such as:

module.weight.data.normal_(0.0, 0.02)

PyTorch’s nn.init functions already operate without autograd tracking. Using the documented functions is clearer and avoids normalizing unsafe .data manipulation as an everyday practice.

Passing the wrong nonlinearity

This is wrong for a standard ReLU layer:

nn.init.kaiming_normal_(
    layer.weight,
    nonlinearity="linear",
)

It produces a different gain from the ReLU configuration.

Likewise, a leaky ReLU layer with slope 0.2 should not be initialized as though its slope were the default value.

Assuming every layer should use the hidden-layer rule

The layer before a ReLU and the final prediction layer have different jobs. Initializing both identically can be defensible in some recipes, but it should be an explicit decision rather than an accident.

Ignoring matrix orientation

PyTorch’s fan calculations assume the conventional linear-layer orientation in which weights have shape (fan_out, fan_in) and are used as x @ weight.T.

For a custom parameter used directly as x @ weight, the stored shape and fan interpretation differ. Initialize the transposed tensor or calculate the fan values consistently with the actual multiplication.

Treating normalization as a complete cure

Batch normalization and layer normalization reduce sensitivity to activation scale, while residual connections provide shorter signal paths. They do not make initialization irrelevant.

Poor initialization can still affect:

• The scale entering the first normalization layer

• Residual-branch magnitude

• Attention logits

• Gating behavior

• Optimizer transients

• Mixed-precision overflow

• The relative learning speed of different branches

Initialization, normalization, residual design, and optimizer settings form a system.

Reproducibility requires more than setting one seed

A fixed seed makes an initialization experiment easier to compare, but exact reproducibility may also depend on:

• Device type

• PyTorch version

• CUDA and accelerator libraries

• Deterministic algorithm settings

• Data-loader worker seeds

• Operation ordering

• Distributed process count

• Floating-point precision

For initialization ablations, use the same model-construction seed for each scheme and regenerate the model before applying the chosen initializer. Otherwise, one method may receive a different sequence of random samples simply because earlier code consumed random numbers.

A clean comparison follows this pattern:

import torch


results = {}

for scheme in ("tiny", "xavier_relu", "he"):
    torch.manual_seed(11)
    model = build_model()
    initialize(model, scheme)
    results[scheme] = run_experiment(model)

For stronger evidence, repeat each configuration across several seeds and report a distribution rather than one best run.

When not to override the defaults

Custom initialization is most valuable when:

• Building a deep network from scratch

• Removing normalization from an architecture

• Introducing unusual nonlinearities

• Changing width or depth substantially

• Adding a new prediction head

• Debugging vanishing or exploding signals

• Reproducing a paper’s exact training recipe

• Designing custom residual or recurrent blocks

Preserve existing initialization when:

• Loading pretrained weights

• Fine-tuning a model with an established recipe

• Using a library architecture whose initializer is part of its design

• You have not yet measured a stability problem

• A checkpoint conversion must remain numerically faithful

“Custom” is not automatically better. Explicit, measured, architecture-aware initialization is better.

A pre-training checklist

Before committing expensive compute, verify:

• The initialization runs exactly once.

• Loaded checkpoint parameters are not overwritten.

• Hidden-layer initialization matches the following activation.

• Output layers are handled deliberately.

• Bias initialization is intentional.

• Forward activations remain finite across depth.

• Activation standard deviations stay within a usable range.

• Early-layer gradients are nonzero and finite.

• Shared parameters are not accidentally initialized multiple times.

• Several random seeds produce similar qualitative behavior.

• Mixed-precision runs do not overflow on the first batch.

• The selected policy is documented beside the model definition.

Then run a small controlled ablation: keep the data, architecture, optimizer, learning rate, and seed fixed, and change only the initialization. Record activation statistics, first-step gradient norms, and early loss curves.

Do that experiment before tuning the optimizer. Your next “optimization problem” may actually be an initialization problem that began—and could have been solved—before step one.