Backpropagation › Addressing Vanishing Gradients in Practice

The vanishing gradient problem shaped much of neural network architecture design over the last decade. Understanding the problem from first principles helps explain why modern architectures look the way they do.

The most influential solution to vanishing gradients is the residual connection, introduced by He et al. (2015) in the ResNet paper. Instead of computing $\mathbf{a}^{(l)} = f(\mathbf{z}^{(l)})$, a residual block computes:

$$\mathbf{a}^{(l)} = f(\mathbf{z}^{(l)}) + \mathbf{a}^{(l-1)}$$

The addition of the input $\mathbf{a}^{(l-1)}$ creates a shortcut path that bypasses the nonlinear transformation. From the backpropagation perspective, the gradient through an addition node is simply the identity: the gradient flows through unchanged. This means there is always a direct gradient highway from the loss all the way back to the earliest layers, regardless of how many layers intervene. The network can learn to use the residual path as heavily as needed, and in the early stages of training when the weights are small, the residual connections dominate and keep gradients flowing.

Transformer models, which power modern language models, use residual connections around both the self-attention sublayer and the feed-forward sublayer in every block. This is not coincidence: without them, even relatively shallow transformers (12-24 layers) would suffer from vanishing gradients and be very difficult to train.

Weight initialization also affects gradient flow. If weights are initialized too small, activations shrink layer by layer in the forward pass, and the gradients are correspondingly small in the backward pass. If weights are too large, activations grow and saturate, again killing gradients.

Xavier (Glorot) initialization sets the variance of weights to $\frac{2}{n_{\text{in}} + n_{\text{out}}}$, which keeps the variance of activations and gradients approximately constant across layers for tanh networks. He initialization sets variance to $\frac{2}{n_{\text{in}}}$, better suited for ReLU networks because ReLU cuts the variance roughly in half (by zeroing negative activations). Getting initialization right is often the difference between a network that starts training smoothly and one that immediately diverges or stalls.

Batch normalization (Ioffe and Szegedy, 2015) normalizes the pre-activation values at each layer to have zero mean and unit variance over the batch. By keeping activations in a reasonable range, it prevents both saturation (which kills gradients) and explosion (which causes instability). It also provides a direct regularization effect and allows higher learning rates, substantially speeding up training. The gradient through a batch normalization layer is more complex than through a plain activation (it involves the batch statistics), but it has the property of preventing activations from drifting into the extreme values where activation derivatives vanish.

Layer normalization, the variant used in transformers, applies the same idea but normalizes across the feature dimension rather than the batch dimension, making it better suited for variable-length sequences and training with batch size 1.