PyTorch Gradient Flow

Gradient flow refers to how gradients travel backwards through a neural network during training. A healthy network passes strong, consistent gradients all the way to its first layer. An unhealthy network suffers from gradients that either vanish to near zero or explode to enormous values — both of which prevent effective learning.

The Chain Rule and Gradient Flow

Backpropagation applies the chain rule of calculus. Gradients at each layer get multiplied together as they travel backwards. If each multiplication slightly reduces the gradient, fifty layers of multiplication can shrink it to essentially zero.

Forward pass (left to right):
Input → Layer1 → Layer2 → Layer3 → Output → Loss

Backward pass (right to left):
Loss → grad3 → grad2 → grad1 → weight updates

Each grad = previous_grad × local_derivative

Vanishing Gradients

Vanishing gradients occur when the local derivatives in a layer are small (less than 1). Each multiplication shrinks the gradient further. By the time it reaches early layers, the gradient is so tiny that those layers barely update — they fail to learn.

The sigmoid activation function was a major cause of vanishing gradients. Its derivative has a maximum value of 0.25. Multiply 0.25 by itself ten times and you get 0.000001 — essentially nothing. ReLU replaced sigmoid in most architectures precisely to solve this problem.

Sigmoid derivative max: 0.25
After 10 layers: 0.25^10 ≈ 0.00000095
                              ↑
                    Gradient nearly dead

Exploding Gradients

The opposite problem: if local derivatives are large (greater than 1), each multiplication amplifies the gradient. After many layers, the gradient becomes astronomically large. The weight update overshoots — the model learns nothing useful and the loss may become NaN.

Gradient value per layer: 2.0
After 10 layers: 2.0^10 = 1024 ← gradient too large
After 20 layers: 2.0^20 = 1,048,576 ← model breaks

Solutions to Vanishing Gradients

ReLU Activation

ReLU (Rectified Linear Unit) returns the input directly for positive values. Its derivative is 1 for all positive inputs — no shrinking. This keeps gradients stable through many layers.

Batch Normalization

Batch normalization normalizes layer inputs during training, keeping activations in a range that avoids saturating activation functions and helps gradients flow more evenly.

Residual Connections (Skip Connections)

Used in ResNet and transformers, residual connections add a direct path from an earlier layer's output to a later layer's input. This gives gradients a highway to travel backwards without passing through many multiplications.

Standard layer:
  input → conv → activation → output

Residual connection:
  input → conv → activation → + output
    │                          ↑
    └──────────────────────────┘ (skip path)

Gradient can flow directly through the skip path — no vanishing risk

Solutions to Exploding Gradients

Gradient Clipping

Gradient clipping caps the gradient norm at a maximum value before updating weights. If the gradient is too large, it gets scaled down to the threshold.

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

Weight Initialization

Starting with weights that are too large causes exploding gradients from the very first step. Proper initialization — Xavier or He initialization — sets initial weights at a scale that supports stable gradient flow.

Visualizing Gradient Flow

You can inspect gradient magnitudes layer by layer to diagnose flow problems.

import torch.nn as nn

def check_gradients(model):
    for name, param in model.named_parameters():
        if param.grad is not None:
            g = param.grad.abs()
            print(f"{name:30s}  mean={g.mean():.6f}  max={g.max():.6f}")
        else:
            print(f"{name:30s}  NO GRADIENT")

Summary

Gradient flow describes how training signals travel backwards through a network. Vanishing gradients occur when layer derivatives are small — early layers stop learning. Exploding gradients occur when layer derivatives are large — the model destabilizes. ReLU activation, batch normalization, and residual connections address vanishing gradients. Gradient clipping and proper weight initialization address exploding gradients. Monitoring gradient magnitudes during training helps you catch these problems early and apply the right fix.