Weight Initialization › Diagnosing Initialization Problems
Language AI HandbookAll chapters Weight Initialization › Diagnosing Initialization ProblemsOne section · Sections in this chapter16 / 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
When a network fails to train, determining whether initialization is the culprit requires specific diagnostic tools. Simply observing poor performance after many training epochs does not distinguish initialization failure from other problems like inappropriate learning rates, overfitting, or data issues.
The most informative diagnostic is to examine activation statistics in the first few forward passes, before any gradient updates. If activations are nearly zero at the output of every layer, you have a vanishing activation problem that is almost certainly initialization-related. If activations are extremely large or if many neurons have saturated activations (tanh outputs close to , sigmoid outputs close to 0 or 1), you have an explosion or saturation problem.
Gradient statistics after the first backward pass provide complementary information. Computing the gradient norm at each layer and plotting it across depth reveals whether gradients are vanishing or exploding as they flow backward through the network. A healthy gradient flow shows gradient norms roughly comparable across all layers. Vanishing gradients appear as norms that decay exponentially with distance from the output; exploding gradients appear as norms that grow exponentially.
def diagnose_initialization(model, input_tensor):
"""
Analyze activation and gradient statistics at initialization.
Returns activation statistics per layer.
"""
model.zero_grad()
activation_stats = []
# Forward pass with hooks to capture activations
hooks = []
activations_captured = []
def hook_fn(module, input, output):
if isinstance(output, torch.Tensor):
activations_captured.append(
{
"mean": output.detach().mean().item(),
"std": output.detach().std().item(),
"frac_dead": (output.detach() == 0).float().mean().item(),
}
)
for name, module in model.named_modules():
if isinstance(module, nn.Linear):
hooks.append(module.register_forward_hook(hook_fn))
output = model(input_tensor)
loss = output.sum()
loss.backward()
for hook in hooks:
hook.remove()
return activations_capturedActivation diagnostics at initialization: ============================================================ zeros: Layer 1: mean=+0.0000, std=0.0000, dead=100.0% Layer 2: mean=+0.0000, std=0.0000, dead=100.0% Layer 3: mean=+0.0000, std=0.0000, dead=100.0% Layer 4: mean=+0.0000, std=0.0000, dead=100.0% he_normal: Layer 1: mean=+0.0295, std=1.4281, dead=0.0% Layer 2: mean=+0.1051, std=1.4485, dead=0.0% Layer 3: mean=+0.0879, std=1.6974, dead=0.0% Layer 4: mean=-1.0546, std=1.4966, dead=0.0% xavier_normal: Layer 1: mean=-0.0113, std=0.9139, dead=0.0% Layer 2: mean=+0.0317, std=0.6534, dead=0.0% Layer 3: mean=+0.0206, std=0.4679, dead=0.0% Layer 4: mean=-0.0132, std=0.3902, dead=0.0%
This diagnostic output immediately shows the difference between initialization schemes. With zero initialization, standard deviations collapse to zero after the first layer and the dead neuron fraction is 100% from the start. With He initialization, standard deviations remain stable across layers and the dead neuron fraction is around 50% (expected for a ReLU network, since the network is applying ReLU to symmetrically distributed pre-activations). With Xavier on a ReLU network, you see a progressive decay in standard deviation, showing the mismatch between Xavier's tanh-oriented design and the ReLU activation.