Vanishing & Exploding Gradients Explorer

🔍 Why Gradients Vanish or Explode

During backpropagation, gradients are computed using the chain rule. For a network with \(L\) layers, the gradient at layer \(l\) involves multiplying gradients from all subsequent layers:

Each term \(\frac{\partial a^{(k+1)}}{\partial a^{(k)}}\) depends on the derivative of the activation function and the weight magnitudes. When these terms are consistently less than 1, their product vanishes exponentially. When greater than 1, it explodes.

The Role of Activation Functions

Activation function derivatives directly impact gradient magnitude. Sigmoid's maximum derivative is 0.25, meaning gradients shrink by at least 75% at each layer. ReLU's derivative is either 0 or 1, preserving gradient magnitude for active neurons.

⚖️ Weight Initialization Strategies

Weight initialization determines the starting point for optimization. Poor initialization can doom training before it begins. The goal is to maintain consistent variance of activations and gradients across layers.

Xavier / Glorot Initialization

Designed for Sigmoid and Tanh activations. Keeps variance stable by accounting for both forward (fan-in) and backward (fan-out) pass:

He / Kaiming Initialization

Designed for ReLU networks. Accounts for the fact that ReLU zeros out half of its inputs, requiring larger initial weights:

LeCun Initialization

Predecessor to Xavier, optimized for SELU activation in self-normalizing networks:

🐍 Inspecting Gradients in PyTorch

Here's a simple example showing how to inspect gradient magnitudes at each layer in a real PyTorch network. This is exactly what this visualization tool simulates.

# Install: pip install torch
import torch
import torch.nn as nn

# Create a simple deep network
class DeepNetwork(nn.Module):
    def __init__(self, depth=10, width=64, activation='relu'):
        super().__init__()
        
        # Choose activation function
        act_fn = {
            'relu': nn.ReLU(),
            'sigmoid': nn.Sigmoid(),
            'tanh': nn.Tanh(),
        }[activation]
        
        # Build layers
        layers = []
        for i in range(depth):
            layers.append(nn.Linear(width, width))
            layers.append(act_fn)
        self.layers = nn.Sequential(*layers)
    
    def forward(self, x):
        return self.layers(x)

# Initialize network
model = DeepNetwork(depth=10, activation='sigmoid')  # Try 'relu' too!

# Create dummy input and target
x = torch.randn(32, 64)  # batch_size=32, features=64
target = torch.randn(32, 64)

# Forward pass
output = model(x)
loss = nn.MSELoss()(output, target)

# Backward pass (computes gradients)
loss.backward()

# Inspect gradient magnitudes at each layer
print("Gradient norms per layer:")
print("-" * 40)

for name, param in model.named_parameters():
    if param.grad is not None:
        grad_norm = param.grad.norm().item()
        print(f"{name:20s} | grad norm: {grad_norm:.6f}")

Example Output (Sigmoid — Vanishing)

Gradient norms per layer:
----------------------------------------
layers.0.weight      | grad norm: 0.000003  ← Vanished!
layers.0.bias        | grad norm: 0.000001
layers.2.weight      | grad norm: 0.000018
layers.2.bias        | grad norm: 0.000006
...
layers.16.weight     | grad norm: 0.089421
layers.16.bias       | grad norm: 0.031245
layers.18.weight     | grad norm: 0.284719  ← Healthy
layers.18.bias       | grad norm: 0.098234

Notice how gradients at early layers (0, 2) are orders of magnitude smaller than later layers (16, 18). This is the vanishing gradient problem in action!

Quick Fix: Use ReLU + He Initialization

# Change activation to ReLU
model = DeepNetwork(depth=10, activation='relu')

# Apply He initialization
for m in model.modules():
    if isinstance(m, nn.Linear):
        nn.init.kaiming_normal_(m.weight, nonlinearity='relu')
        nn.init.zeros_(m.bias)

With ReLU and He initialization, gradient norms stay much more consistent across layers, enabling effective training of deep networks.

📖 Glossary of Terms

Vanishing Gradients

A phenomenon where gradients become exponentially smaller as they propagate backward through layers, causing early layers to receive near-zero updates and stop learning effectively.

Exploding Gradients

A phenomenon where gradients become exponentially larger during backpropagation, leading to unstable training, huge weight updates, and NaN values in the loss.

Saturation

When activation function inputs fall into regions where the derivative is near zero (e.g., very large or small inputs for Sigmoid), causing gradients to vanish.

Dead Neurons

Neurons that always output zero because they're stuck in ReLU's inactive region (negative inputs). Once dead, they receive zero gradients and never recover.

Fan-in / Fan-out

Fan-in is the number of input connections to a neuron; fan-out is the number of output connections. Used to calculate proper weight initialization variance.

Gradient Clipping

A technique to prevent exploding gradients by capping gradient magnitudes to a maximum threshold during backpropagation.

❓ Frequently Asked Questions

📚 Recommended Reading