Diagnosing Vanishing Gradients

Run a deep MLP forward and backward and measure the gradient norm reaching each layer.

Implement layer_grad_norms(depth, width, activation, seed=0):

• weights W_l drawn with rng.normal(0, 1/sqrt(width), (width, width)) for each layer in order, from one np.random.default_rng(seed), then an input x = rng.normal(size=width)
• forward: h_l = act(W_l h_{l-1}) for activation in "sigmoid", "tanh", "relu"
• backward from upstream gradient ones(width) at the output, through each activation and weight matrix
• return a list of the gradient norms with respect to each layer's pre-activation, from the first layer to the last