Weight Initialization and Gradient Flow in Deep Networks

Backpropagation through a deep network multiplies layer derivatives. Depending on their sizes and directions, a gradient can shrink, grow, or remain useful as it travels backward. Initialization sets the starting scales of these transformations. Xavier and He initialization aim to control typical signal sizes under simplifying assumptions; they do not guarantee successful training.

Why depth multiplies the problem

Consider an \(L\)-layer linear network with zero biases and \(W^{[l]}=cI\), where \(I\) is the identity matrix. Its output is \(c^Lx\). For \(L=50\), the scaling is about 117 when \(c=1.1\) and 0.00515 when \(c=0.9\). Given a fixed gradient at the output, the gradient with respect to the input is multiplied by the same \(c^L\). Parameter gradients also involve the activations entering each layer and are not all scaled by that single factor.

For a layer \(a=\phi(Wx+b)\), backpropagation computes \(g_x=W^T(g_a\odot\phi'(Wx+b))\): the incoming gradient \(g_a\) is multiplied elementwise by activation derivatives, then by the transposed weights. Matrix directions matter as well as scale. Activation statistics are useful diagnostics, but they do not determine gradient flow on their own.

import numpy as np

def signal_probe(scale, n=256, layers=12, seed=0):
    rng = np.random.default_rng(seed)
    a = rng.standard_normal((n, 512))
    cache, stds = [], []
    for _ in range(layers):
        W = rng.standard_normal((n, n)) * scale
        a = np.tanh(W @ a)
        cache.append((W, a))
        stds.append(a.std())
    g = rng.standard_normal(a.shape)
    output_rms = np.sqrt(np.mean(g*g))
    for W, a in reversed(cache):
        g = W.T @ (g * (1-a*a))
    gain = np.sqrt(np.mean(g*g)) / output_rms
    return stds[0], stds[-1], gain

print("scale first_std last_std input/output_gradient_rms")
for scale in (0.01, 1.0, np.sqrt(1/256)):
    first, last, gain = signal_probe(scale)
    print(f"{scale:.4f} {first:.4f} {last:.4f} {gain:.3e}")
# scale first_std last_std input/output_gradient_rms
# 0.0100 0.1564 0.0000 2.826e-10
# 1.0000 0.9745 0.9740 4.352e+05
# 0.0625 0.6279 0.2119 2.476e-01

The arrays have 256 units in rows and 512 examples in columns. Each run resets the seed so the input, underlying random weight draws, and injected output gradient match across scales. The final column is the root-mean-square (RMS) input gradient divided by the injected output-gradient RMS. It probes one backward direction without fitting a task or computing weight gradients. Read it alongside the activation standard deviations; neither statistic proves that a model will train.

At scale 0.01 the input gradient shrinks to about \(2.8\times10^{-10}\) of the output gradient; the printed final activation standard deviation rounds to zero. At scale 1, tanh outputs are near saturation, yet this sampled backward signal grows by about \(4.4\times10^5\). Large weights and the remaining unsaturated paths can outweigh the small local activation derivatives. At scale 0.0625, activation standard deviation declines from 0.6279 to 0.2119 and the gradient RMS ratio is about 0.248. This is a milder change in this experiment, not exact signal preservation.

The L-layer forward and backward passes are written out in Deep Neural Networks: Architecture and Backpropagation.

Symmetry between hidden units

Hidden units remain identical under deterministic, symmetry-preserving updates when their incoming weights and biases match and their outgoing connections also treat them identically. Initializing an ordinary multilayer network entirely to zero creates this problem; with zero-initialized tanh hidden layers, their parameter gradients also vanish. Identical incoming weights in just one layer do not alone guarantee identical future updates, because unequal outgoing weights can break the symmetry.

Zero biases are common when random weights already distinguish the units. Zero weights are valid in some settings: logistic regression has no interchangeable hidden units, and its gradient can be nonzero at zero initialization. Specialized residual branches also use zero initialization deliberately. The concern is preserving unwanted symmetry or blocking gradients, not the number zero by itself.

Xavier initialization for symmetric activations

For a dense weight matrix of shape \((n_{out},n_{in})\), fan-in \(n_{in}\) is the number of inputs to a unit and fan-out \(n_{out}\) is the number of output units. With zero bias, take \(z=\sum_iw_ix_i\) with independent, zero-mean weights of common variance, independent of the inputs. Averaging over these draws gives \(\operatorname{Var}(z)=n_{in}\operatorname{Var}(w)\mathbb E[x_i^2]\) when inputs have the same second moment. If they are centered, that second moment equals their variance. Preserving forward variance then suggests \(\operatorname{Var}(w)=1/n_{in}\). A similar backward approximation, treating incoming gradients as independent of the weights, suggests \(1/n_{out}\). These independence assumptions need not hold exactly in a realized network. Xavier uses a compromise:

\[\text{Var}(w)=\frac{2}{n_{in}+n_{out}}\]

This approximation treats the activation as nearly linear with slope one, as tanh is near zero. It does not account for strong tanh saturation. For \(n_{in}=n_{out}=256\), the variance is \(1/256\) and the normal-distribution standard deviation is \(1/16=0.0625\). The square root matters because the code multiplies standard-normal samples by a standard deviation, not a variance.

He initialization for ReLU

For a symmetric pre-activation \(z\), ReLU gives \(a=\max(0,z)\). Its second moment satisfies \(\mathbb E[a^2]=\tfrac12\mathbb E[z^2]\). This is not the variance: \(\operatorname{Var}(a)=\mathbb E[a^2]-\mathbb E[a]^2\), and ReLU usually has positive mean. For standard-normal \(z\), the second moment is 0.5 but the variance is \(0.5-1/(2\pi)\approx0.341\). Combining the second-moment factor with the weighted-sum calculation motivates forward-preserving He initialization:

\[\text{Var}(w)=\frac{2}{n_{in}}\]

For equal-width layers, this variance is twice Xavier’s, so the sampling standard deviation is \(\sqrt2\) times larger. When fan-in and fan-out differ, their ratio is not generally two. He fan-in targets the forward signal; a fan-out variant targets backward scale. For Leaky ReLU with negative slope \(r\), the analogous variance is \(2/[(1+r^2)n_{in}]\). GELU does not follow this exact rectifier calculation. Xavier is a starting choice for centered, roughly linear activations; sigmoid is not centered and its derivative is at most 0.25, so Xavier does not remove saturation or deep-gradient difficulties. Match initialization to the architecture and activation instead of applying one rule to every layer.

def he_normal(n_out, n_in, rng):
    return rng.standard_normal((n_out, n_in)) * np.sqrt(2.0 / n_in)

def xavier_uniform(n_out, n_in, rng):
    limit = np.sqrt(6.0 / (n_in + n_out))
    return rng.uniform(-limit, limit, size=(n_out, n_in))

def init_parameters(layer_dims, rng, modes):
    if len(modes) != len(layer_dims)-1:
        raise ValueError("Specify an initialization for each weight matrix")
    if any(d <= 0 for d in layer_dims):
        raise ValueError("Layer widths must be positive")
    initializers = {"he": he_normal, "xavier": xavier_uniform}
    params = {}
    for l, mode in enumerate(modes, 1):
        if mode not in initializers:
            raise ValueError(f"Unknown initialization: {mode}")
        n_in, n_out = layer_dims[l-1], layer_dims[l]
        params[f"W{l}"] = initializers[mode](n_out, n_in, rng)
        params[f"b{l}"] = np.zeros((n_out, 1))
    return params

params = init_parameters([784, 256, 128, 10], np.random.default_rng(1),
                         modes=["he", "he", "xavier"])
for k, v in params.items():
    print(k, v.shape, round(float(v.std()), 4))
# W1 (256, 784) 0.0504
# b1 (256, 1) 0.0
# W2 (128, 256) 0.0885
# b2 (128, 1) 0.0
# W3 (10, 128) 0.1223
# b3 (10, 1) 0.0

The example assumes two ReLU hidden layers and an affine output producing ten logits, so it uses He for the hidden layers and Xavier for the output as a starting choice. The uniform limit \(\sqrt{6/(n_{in}+n_{out})}\) follows from \(\operatorname{Var}(U[-a,a])=a^2/3\). Printed standard deviations are sample measurements and need not equal the theoretical values exactly.

Diagnosing which failure you have

Gradient norms help localize a problem but are not a diagnosis. A larger matrix can have a larger norm simply because it has more entries. Report both the norm and RMS per entry, inspect zero or nonfinite gradients, and compare each layer across steps. Also track actual update size relative to weight size where weights are nonzero. There is no universal healthy first-to-last ratio, and rounding tiny values to six decimals can conceal them.

def gradient_stats(grads, L):
    for l in range(1, L+1):
        g = grads[f"dW{l}"]
        norm = np.linalg.norm(g)
        rms = norm / np.sqrt(g.size)
        print(f"layer {l}: norm={norm:.3e} rms={rms:.3e}")

toy_grads = {"dW1": np.full((100, 100), 1e-4),
             "dW2": np.full((10, 10), 1e-4)}
gradient_stats(toy_grads, 2)
# layer 1: norm=1.000e-02 rms=1.000e-04
# layer 2: norm=1.000e-03 rms=1.000e-04

The toy matrices have identical gradients per entry, yet their norms differ tenfold because the first has 100 times as many entries. RMS makes that size effect visible; it still does not establish whether either layer is learning effectively.

Gradient clipping

Global-norm clipping bounds the concatenated gradient. For \(\tau>0\), it rescales a nonzero gradient as \(g\leftarrow g\min(1,\tau/\|g\|)\), leaving a zero gradient unchanged. For plain gradient descent this bounds the update norm by \(\eta\tau\). Momentum, adaptive scaling, and separate weight decay can change the final update, so the same bound need not apply to them. The positive rescaling preserves gradient direction; it does not guarantee that a finite step decreases the loss.

Clipping can reduce the effect of large finite gradient spikes. Frequent clipping is a reason to inspect gradient scales, the clipping threshold, data, loss reduction, and optimizer settings; it does not prove that initialization or the learning rate is wrong. Clipping cannot recover gradients that have already become NaN or infinite.

def clip_by_global_norm(grads, max_norm=1.0):
    if not grads:
        raise ValueError("Provide at least one gradient array")
    if not np.isfinite(max_norm) or max_norm <= 0:
        raise ValueError("max_norm must be positive and finite")
    if any(not np.isfinite(g).all() for g in grads.values()):
        raise ValueError("Nonfinite gradient: diagnose before clipping")
    flat = np.concatenate([np.asarray(g, dtype=float).ravel()
                           for g in grads.values()])
    total = float(np.hypot.reduce(flat, initial=0.0))
    if not np.isfinite(total):
        raise ValueError("Gradient norm exceeds the numeric range")
    if total <= max_norm:
        return grads, total
    factor = max_norm / total
    return {k: g * factor for k, g in grads.items()}, total

clipped, original_norm = clip_by_global_norm({"w": np.array([3., 4.])})
print("original norm", original_norm)
print("clipped", clipped["w"])
# original norm 5.0
# clipped [0.6 0.8]

The normalization layers referred to here are compared in Normalization Layers Explained: BatchNorm, LayerNorm, GroupNorm, RMSNorm.

What initialization cannot fix

Initialization targets a useful starting scale under assumptions about inputs and random weights. Neither exact variance preservation in a finite network nor stability after many updates follows automatically. Normalization and residual connections provide additional ways to manage signal propagation, but their effects depend on placement and parameterization.

For a residual block \(y=x+F(x)\), the input gradient is \(g_x=g_y+J_F(x)^Tg_y\), where \(J_F\) is the derivative matrix of the branch. The direct term can help gradients propagate, but the branch can cancel or amplify it. Normalization controls selected statistics, not every gradient direction. Inspect the actual architecture, input scale, loss, and optimization together.

Exercises

1. Depth compounds. A network applies the same scalar factor \(c\) at every layer. Tabulate the total scaling for \(c\in\{0.9,1.0,1.1\}\) at depths 10, 30, and 50, and find the depth at which \(c=1.1\) first exceeds \(100\times\).

You should get: three columns spanning several orders of magnitude, and a threshold depth just under 50.

for L in (10, 30, 50):
    print(L, [f"{c**L:.3g}" for c in (0.9, 1.0, 1.1)])
# 10 ['0.349', '1', '2.59']
# 30 ['0.0424', '1', '17.4']
# 50 ['0.00515', '1', '117']
print(next(L for L in range(1, 200) if 1.1**L > 100))   # 49

2. He versus Xavier under ReLU. Run the activation-standard-deviation experiment with ReLU over 20 layers, comparing scale \(\sqrt{1/n}\) with \(\sqrt{2/n}\). Report the final standard deviation and RMS for each, using matching input and weight draws. Explain which moment the factor of two targets.

You should get: a much smaller forward signal with Xavier in this run; this is not a training experiment.

import numpy as np

def run(scale, n=256, layers=20, seed=0):
    rng = np.random.default_rng(seed)
    a = rng.standard_normal((n, 512))
    for _ in range(layers):
        W = rng.standard_normal((n, n)) * scale
        a = np.maximum(W @ a, 0)
    return a.std(), np.sqrt(np.mean(a*a))

for name, scale in (("Xavier", np.sqrt(1/256)), ("He", np.sqrt(2/256))):
    std, rms = run(scale)
    print(name, "std", round(std, 6), "rms", round(rms, 6))
# Xavier std 0.00069 rms 0.000849
# He std 0.70692 rms 0.869271

3. Clipping is not free. Build a gradient dictionary with one entry of norm 100 and three of norm 0.1, then apply global-norm clipping at \(\tau=1\). Report each entry’s norm afterwards. What happened to the small gradients, and how does per-element clipping change the direction?

You should get: every entry scaled by the same factor, including the ones that were already small.

import numpy as np
grads = {"big": np.full(10, 100/np.sqrt(10)),
         **{f"small{i}": np.full(10, 0.1/np.sqrt(10)) for i in range(3)}}
clipped, total = clip_by_global_norm(grads, max_norm=1.0)
for k, g in grads.items():
    print(k, round(float(np.linalg.norm(g)), 4), "->",
          round(float(np.linalg.norm(clipped[k])), 6))
print("total norm", round(total, 4), "factor", round(1/total, 6))
# big 100.0 -> 0.999999
# small0 0.1 -> 0.001
# small1 0.1 -> 0.001
# small2 0.1 -> 0.001
# total norm 100.0001 factor 0.01

References

• Bengio, Simard, and Frasconi (1994). Learning Long-Term Dependencies with Gradient Descent Is Difficult. IEEE Transactions on Neural Networks.
• Glorot and Bengio (2010). Understanding the Difficulty of Training Deep Feedforward Neural Networks. AISTATS.
• Pascanu, Mikolov, and Bengio (2013). On the Difficulty of Training Recurrent Neural Networks. ICML.
• He, Zhang, Ren, and Sun (2015). Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification. ICCV.