Orthogonality and weight initialization

When a network produces NaNs on the first step, the cause is often its starting weights. Each layer multiplies the signal by a matrix. If the matrix stretches too much, the signal explodes. If it stretches too little, the signal vanishes. Orthogonality and variance explain why a good starting scale is roughly 1 over the square root of the layer's width.

The big idea

Two vectors are orthogonal when their dot product is zero. They are perpendicular, so they point in independent directions. An orthonormal set has unit lengths and is mutually orthogonal. An orthogonal matrix only rotates or reflects. It keeps lengths and angles unchanged.

Random weights do not behave this way. Their effect on the signal size depends on their scale. Stack layers and small per-layer errors compound.

Walking through it: the variance rule

Take one linear layer with no bias. The inputs xᵢ have variance 1. The weights wᵢ are random with variance σ². The output is y = Σ wᵢ xᵢ.

Because the terms are independent with mean zero, their variances add:

Var(y) = d_in · σ²

To keep Var(y) near 1, choose σ ≈ 1/√d_in. For a layer with 784 inputs, σ ≈ 1/28 ≈ 0.036.

After a ReLU, half the units output zero, so the signal loses half its variance per layer. He initialization compensates with σ = √(2 / d_in). For 784 inputs that is about 0.05.

Now compound it. If each layer multiplies the signal size by 1.5, ten layers give 1.5¹⁰ ≈ 58 times the input. If each layer multiplies by 0.5, ten layers give 0.5¹⁰ ≈ 0.001. Both are failures, and the cause is the scale, not the data.

This sketch stacks ten linear layers and prints the spread of the output. It runs as written.

import numpy as np

rng = np.random.default_rng(0)
d = 784
x = rng.normal(size=(64, d))           # unit-variance inputs

for name, scale in [("scale 1.0", 1.0), ("1/sqrt(d)", 1 / np.sqrt(d))]:
    h = x
    for _ in range(10):
        W = rng.normal(size=(d, d)) * scale
        h = h @ W                      # linear layers only, no activation
    print(f"{name:10s} std after 10 layers = {h.std():.3g}")

Output:

scale 1.0  std after 10 layers = 2.96e+14
1/sqrt(d)  std after 10 layers = 1

With scale 1.0, each layer multiplies the spread by about √784 = 28, so ten layers blow up. With 1/√d, the spread stays near 1.

Common mistakes

• Assuming random initialization is always safe. The right scale depends on layer width.
• Using one rule for every activation. Tanh-style layers use the 1/√d rule. ReLU layers use He initialization.
• Treating orthogonality as a cure-all. Orthogonal matrices keep the size of a signal. They do not fix a bad learning rate or bad data.

Check yourself

1. Are [1, 0] and [0, 5] orthogonal? Answer: Yes. Their dot product is 1·0 + 0·5 = 0.
2. A layer has 256 inputs. What standard deviation keeps the output variance near 1? Answer: 1/√256 = 1/16 = 0.0625.

Key takeaways

• Orthogonal vectors point in independent directions, and orthogonal matrices preserve lengths.
• Output variance scales with d_in · σ², so the weight scale should be about 1/√d_in.
• ReLU layers need a larger scale, which He initialization provides.
• If a deep net outputs NaNs on step one, check the weight scale first.

Further learning

Back to the topic overview: linear algebra for ML. The training page covers initialization in the context of a full training run: training neural networks well.