Vanishing and exploding gradients: why a deep, badly-initialized net fails to train

The problem, from first principles

[02-layers/04-weight-initialization] derived Xavier and Kaiming initialization from the requirement "keep activation variance roughly constant across layers." This question demonstrates, numerically, exactly what goes wrong if you DON'T: [02-layers/02-linear-backward]'s linear_backward shows that backpropagating a gradient through one layer involves multiplying by that layer's weight matrix. Chain many layers together, and the gradient reaching an early layer has been multiplied by EVERY layer's weight matrix in between, one after another. If each individual multiplication happens to shrink the gradient's magnitude even slightly (say, by a factor of 0.9), the CUMULATIVE effect across, say, 50 layers is 0.9^50, an almost unimaginably tiny number, meaning early layers receive a gradient signal so close to zero that gradient descent essentially can't move them at all: they stop learning. The mirror-image failure, if each multiplication instead slightly GROWS the gradient (a factor of 1.1), compounds to 1.1^50, an enormous number, which manifests as training instability, nan losses, or wildly oscillating weights (exactly what [03-dl-training/01-optimizers/07-gradient-clipping]'s clip_grad_norm was built to guard against).

This is a fundamentally MULTIPLICATIVE, compounding effect, and that's precisely why it becomes catastrophic specifically for DEEP networks: a shallow, 2-3 layer network barely notices a per-layer factor of 0.9 or 1.1, but a 50-layer network experiences an effect that's roughly factor^50, astronomically far from factor^1.

From theory to code

Implement scalar_gradient_chain(depth, layer_scale), the simplest possible model: layer_scale raised to the depth-th power, representing what happens if every layer scales the gradient by the exact same fixed factor. Then implement matrix_gradient_norms(depth, dim, weight_std, rng), a more realistic simulation: start with a gradient vector of all 1s, and at each of depth layers, multiply it by a FRESH random (dim, dim) weight matrix with entries drawn from N(0, weight_std^2) (exactly the kind of matrix [02-layers/04-weight-initialization]'s formulas are designed to scale correctly), tracking the gradient vector's norm after every layer.

Constraints

• scalar_gradient_chain computes layer_scale ** depth exactly.
• matrix_gradient_norms must return a list of length depth + 1: the starting norm (before any layers), followed by the norm after each of the depth layers.
• Each layer in matrix_gradient_norms must use a FRESH random weight matrix (drawn via the passed-in rng), not the same matrix reused across every layer.
• matrix_gradient_norms must genuinely demonstrate BOTH failure modes: a small weight_std should cause the norm to shrink toward zero across depth, and a large weight_std should cause it to grow explosively.

Hints