What is the vanishing/exploding gradient problem and how do you detect and fix it in PyTorch?

During backpropagation, gradients are computed via repeated multiplication through the chain rule. In deep networks, this can cause gradients to shrink toward zero (vanishing) or grow toward infinity (exploding) as they propagate backward through many layers, preventing effective training.

import torch
import torch.nn as nn

model     = nn.LSTM(input_size=10, hidden_size=128, num_layers=3, batch_first=True)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

x = torch.randn(32, 20, 10)
output, _ = model(x)
loss = output.sum()

optimizer.zero_grad()
loss.backward()

# ── Detect: monitor gradient norms
total_norm = 0.0
for p in model.parameters():
    if p.grad is not None:
        total_norm += p.grad.data.norm(2).item() ** 2
total_norm = total_norm ** 0.5
print(f"Gradient norm: {total_norm:.4f}")
# Very small (~1e-6) → vanishing; very large (~1e3+) → exploding

# ── Fix 1: Gradient clipping — caps the gradient norm before the step
nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()

# ── Fix 2: Better weight initialisation (He init for ReLU networks)
def init_weights(m):
    if isinstance(m, nn.Linear):
        nn.init.kaiming_uniform_(m.weight, nonlinearity="relu")
        nn.init.zeros_(m.bias)

mlp = nn.Sequential(nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 64))
mlp.apply(init_weights)

# ── Fix 3: Batch Normalization — stabilises layer input distributions
class StableNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(64, 64), nn.BatchNorm1d(64), nn.ReLU(),
            nn.Linear(64, 64), nn.BatchNorm1d(64), nn.ReLU(),
        )
    def forward(self, x):
        return self.net(x)

# ── Fix 4: Residual / skip connections — gradient highway
class ResidualBlock(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(dim, dim), nn.ReLU(), nn.Linear(dim, dim))
    def forward(self, x):
        return x + self.net(x)   # gradient flows through x directly

Invest now in Acorns!!! 🚀
Join Acorns and get your $5 bonus!

Acorns is a micro-investing app that automatically invests your "spare change" from daily purchases into diversified, expert-built portfolios of ETFs. It is designed for beginners, allowing you to start investing with as little as $5. The service automates saving and investing. Disclosure: I may receive a referral bonus.

Invest now!!! Get Free equity stock (US, UK only)!

Use Robinhood app to invest in stocks. It is safe and secure. Use the Referral link to claim your free stock when you sign up!.

The Robinhood app makes it easy to trade stocks, crypto and more.

Webull! Receive free stock by signing up using the link: Webull signup.