What this usually means
Numerical instability in backpropagation: gradients grow exponentially due to deep architectures, improper weight initialization, high learning rates, or unstable loss functions (e.g., log(0) in cross-entropy). The explosion causes weights to become Inf/NaN, ruining training. It often appears at early steps or after a batch with extreme activations.
The first ten minutes — establish facts before touching code.
- 1Insert torch.autograd.set_detect_anomaly(True) and rerun; it pinpoints the first NaN backward op.
- 2Log gradient norms every step: print(torch.nn.utils.clip_grad_norm_(model.parameters(), float('inf'))).
- 3Check input data for NaNs/Infs: torch.isnan(input).any() and torch.isinf(input).any().
- 4Inspect loss function: ensure no log(0) or division by small numbers (e.g., add epsilon).
- 5Reduce learning rate by 10x; if NaN persists, the issue is likely initialization or architecture.
The specific files, logs, configs, and dashboards that usually own this bug.
- searchTraining loop: gradient clipping call (clip_grad_norm_) and optimizer step.
- searchForward pass: output of loss function (e.g., cross_entropy, mse).
- searchModel __init__: weight initialization (e.g., nn.init.xavier_uniform_).
- searchData pipeline: batch that triggers NaN (log it).
- searchGradient histogram: torch.nn.utils.clip_grad_norm_ with inf norm to see unscaled norms.
- searchLogs: check for 'NaN' or 'inf' in loss values; look at step where gradient norm spikes.
- searchCheckpoint: save model before NaN occurs and compare weights with after.
Practical causes, not theory. These are the things you will actually find.
- warningLearning rate too high, overshooting loss surface.
- warningWeights initialized too large (e.g., normal(0,1) instead of small uniform).
- warningNo gradient clipping on recurrent or deep networks.
- warningLoss function with unbounded outputs (e.g., MSE on unbounded targets).
- warningInput data not normalized (huge values saturate sigmoid/tanh).
- warningMixed precision training without gradient scaling (AMP).
- warningBatch size too small causing high variance gradients.
Concrete fix directions. Pick the one that matches your root cause.
- buildApply gradient clipping: torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0).
- buildReduce learning rate: schedule or reduce to 1e-4 and gradually increase.
- buildUse proper weight initialization: xavier_uniform_ for tanh, kaiming_uniform_ for ReLU.
- buildNormalize inputs to zero mean unit variance; for images, scale to [0,1] or standardize.
- buildAdd epsilon to loss denominators (e.g., log(p + 1e-8) in cross-entropy).
- buildUse torch.nn.utils.clip_grad_value_ to cap gradient values directly.
- buildSwitch to Adam optimizer (adaptive LR) or use learning rate warmup.
A fix you cannot prove is a guess. Close the loop.
- verifiedRun training with gradient clipping; loss stabilizes and converges.
- verifiedLog gradient norms every 10 steps; they stay below 1e2.
- verifiedCheck model weights after 100 steps: no NaN/Inf.
- verifiedUse a sanity check: overfit a single batch (if loss goes to zero, gradients are fine).
- verifiedMonitor loss curve: no spikes, smooth descent.
- verifiedValidate on held-out set: accuracy increases consistently.
Things that make this bug worse or harder to find.
- warningClipping too aggressively (e.g., max_norm=0.01) kills gradient signal.
- warningDisabling anomaly detection after fixing; it's a cheap safety net.
- warningIgnoring data pipeline: a single NaN pixel can explode gradients.
- warningUsing gradient clipping as a band-aid without addressing root initialization/architecture.
- warningNot normalizing inputs when using activation functions like sigmoid/tanh.
- warningSetting too high a learning rate after clipping; clipping doesn't fix the underlying instability.
ResNet-50 on CIFAR-10: Loss goes NaN at step 42
Timeline
- 09:00Start training ResNet-50 from scratch with random init.
- 09:02Step 15: loss = 2.3, gradient norm = 12.4.
- 09:03Step 42: loss = NaN, gradient norm = inf.
- 09:05Check model weights: all NaN in final linear layer.
- 09:10Enable anomaly detection; backward on cross_entropy throws error.
- 09:15Inspect logits before loss: max logit = 1e4, softmax underflows.
- 09:20Reduce learning rate to 1e-4; NaN still appears at step 100.
- 09:30Add gradient clipping max_norm=10; training stabilizes.
- 09:45Switch to Kaiming init; remove clipping, still stable.
I started training a ResNet-50 on CIFAR-10 with random initialization (normal(0,1)) and Adam at 1e-3. By step 42, loss went NaN. The gradient norm was inf, and model weights became NaN. I suspected gradient explosion.
I enabled anomaly detection, which pointed to the cross-entropy backward step. Inspecting logits revealed values up to 1e4, causing softmax to output extreme probabilities that underflow to zero, making log(0) produce NaN. The root cause was the large initial weights creating huge activations.
I first reduced learning rate, but NaN persisted. Applying gradient clipping with max_norm=10 stabilized training immediately. Later, I replaced the init with Kaiming uniform (recommended for ReLU) and removed clipping—training remained stable. The fix was proper initialization, with clipping as a safety net.
Root cause
Weight initialization with normal(0,1) produced large weights, causing activations to explode, leading to extreme softmax outputs and log(0) in cross-entropy loss.
The fix
Changed initialization to nn.init.kaiming_uniform_(layer.weight) for all Conv2d and Linear layers. Also applied gradient clipping max_norm=10 as a temporary measure.
The lesson
Always use appropriate weight initialization (Kaiming for ReLU, Xavier for tanh). Gradient clipping is a crutch; fix the root cause first.
In deep networks, gradients multiply through layers (chain rule). If weights are large (>1), gradients can grow exponentially, causing numerical overflow to Inf/NaN. Common triggers: unbounded activation (ReLU) with large weights, recurrent networks (vanishing/exploding), or loss functions that produce large gradients (e.g., MSE with outliers).
Example: In a 50-layer ResNet, if each layer multiplies by 1.1, the gradient can grow by 1.1^50 ≈ 117. With weight norm >1, explosion is guaranteed. The first sign is a sudden spike in gradient norm, followed by NaN loss.
Run with torch.autograd.set_detect_anomaly(True) to get a traceback on the first NaN backward op. It slows training but pinpoints the exact layer. Use it on a small subset.
Log gradient norms every step: total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), float('inf')). This gives the L2 norm without clipping. Watch for jumps >1e4.
Check model parameters for NaN: any(torch.isnan(p).any() for p in model.parameters()). Insert after optimizer.step().
Inspect activations: register forward hooks on suspicious layers to see mean/std of outputs. Large pre-activations (>1e3) indicate explosion.
Clip by norm: torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm). Common values: 1.0 for LSTM, 10.0 for CNNs. Too low (0.01) slows training; too high (1e3) is useless. Start with 1.0 and adjust based on gradient norm logs.
Clip by value: torch.nn.utils.clip_grad_value_(model.parameters(), clip_value=0.5). This caps each gradient component, preventing any single gradient from being huge. Useful when only a few gradients explode.
For mixed precision (AMP), use GradScaler to scale loss before backward, then unscale before clipping. Forgetting to unscale leads to under-clipping.
For ReLU: use nn.init.kaiming_uniform_(weight, mode='fan_in', nonlinearity='relu'). This keeps activation variance ~1. For tanh: use xavier_uniform_. Avoid normal(0,1) unless the network is very shallow.
Batch normalization helps by rescaling layer outputs, but if weights are too large, BN can't compensate. Still, initialize properly.
Test initialization: run a forward pass on a random batch and check output variance. If it's >1e3, the init is too large.
Cross-entropy with logits: use torch.nn.CrossEntropyLoss which internally applies log-softmax without underflow. Avoid manual log(softmax(x)).
If implementing custom loss, add a small epsilon: log(x + 1e-8) to prevent log(0). Similarly, for division, add epsilon to denominator.
Check for extreme logits: if logits exceed 1e4, softmax underflows. Use temperature scaling (divide logits by T>1) to soften probabilities.
Frequently asked questions
Why does NaN loss only appear at certain batches?
Gradient explosion often triggers on specific batches with edge-case data (e.g., outliers, corrupted images). Those batches produce larger activations, which amplify gradient magnitudes. Running with a fixed seed can reproduce the exact batch; inspect it for anomalies.
Does gradient clipping always fix NaN loss?
No. Clipping masks the symptom but doesn't fix the underlying instability (bad initialization, architecture, or learning rate). It should be combined with proper initialization and learning rate scheduling. If clipping is set too high, explosion still happens; too low, training stalls.
How do I differentiate gradient explosion from vanishing gradients?
Explosion: gradient norm >1e4, loss goes NaN or large spike. Vanishing: gradient norm <1e-8, loss plateaus. Check gradient norms: explosion has inf/near-inf values; vanishing has near-zero. Both can occur in different layers of the same network.
Can mixed precision training cause NaN loss?
Yes. In float16, gradients can underflow/overflow more easily. Use loss scaling (GradScaler) and dynamic scaling. If you see NaN with AMP, disable AMP to see if it's the cause. Also ensure all layers are compatible with float16 (e.g., layernorm in float32).
Should I clip gradients before or after optimizer step?
Clip after backward() and before optimizer.step(). The order: loss.backward() -> clip_grad_norm_ -> optimizer.step(). Clipping modifies gradients in-place; optimizer then uses the clipped gradients. If you clip after step, it has no effect.