Initialisation and variance
The variance of a sum of independent terms is the sum of their variances. A unit computes w @ x over n inputs; if each weight has variance v and each input variance 1, the output has variance n v. Choose v = 1 / n and the signal keeps its size from layer to layer. Too small, and it shrinks to nothing after a few layers; too large, and it explodes.
ReLU zeroes half of its inputs, halving the second moment, so ReLU networks use v = 2 / n, which is He initialisation, standard deviation sqrt(2 / n). Xavier initialisation, sqrt(2 / (n_in + n_out)), balances the forward and backward passes for tanh.
import numpy as np
rng = np.random.default_rng(0)
x = rng.normal(size=(1000, 256))
for std in (0.01, np.sqrt(2 / 256)):
h = x
for _ in range(10):
W = rng.normal(0, std, (256, 256))
h = np.maximum(h @ W, 0)
print(round(std, 3), h.std().round(4))Frameworks pick these scales for you, but a custom layer or a very deep stack is where knowing them pays off.
Your task
Write he_std(fan_in) and xavier_std(fan_in, fan_out) as floats. Write relu_stack_stds(x, width, depth, std_fn, seed): starting from the batch x, apply depth layers h = relu(h @ W) with W of shape (h.shape[1], width) drawn from np.random.default_rng(seed) normal with standard deviation std_fn(fan_in), and return the list of each layer's pre-activation standard deviations (of h @ W).
The explanation opens here once every check passes.