Your Loss Is NaN. Now What?

The Dreaded Loss: NaN Exception

You start a 10-hour deep learning training job. At epoch 3, the loss output suddenly displays:

Epoch 1: Loss = 2.412
Epoch 2: Loss = 1.845
Epoch 3: Loss = nan  <── TRAINING IS BROKEN!

Once loss becomes NaN (Not a Number), backpropagation propagates NaN gradients to all model weights, permanently corrupting the model.

Here is the 5-Step Diagnostic Checklist to fix NaN loss.

Loss is NaN ──► 1. Check Inputs ──► 2. Clip Gradients ──► 3. Fix Log/Div ──► 4. Lower LR ──► 5. Enable Anomaly Detection

The 5 Step NaN Loss Diagnostic Checklist

┌──────────────────────────┬──────────────────────────┬──────────────────────────┐
│ 1. CHECK INPUT FEATURES  │ 2. CLIP GRADIENTS        │ 3. SAFE LOG & DIVISION   │
├──────────────────────────┼──────────────────────────┼──────────────────────────┤
│ Check datasets for       │ Cap maximum gradient     │ Add tiny epsilon offset  │
│ unhandled NaNs, nulls,   │ norm to prevent exploding│ e.g. torch.log(p + 1e-12)│
│ or un-scaled huge values!│ weights!                 │ to avoid log(0) = -inf!  │
└──────────────────────────┴──────────────────────────┴──────────────────────────┘

┌──────────────────────────┬──────────────────────────┐
│ 4. LOWER LEARNING RATE   │ 5. AUTOGRAD ANOMALY DETECT│
├──────────────────────────┼──────────────────────────┤
│ Reduce learning rate by  │ Enable PyTorch anomaly   │
│ 10x or add warm-up       │ detection to identify    │
│ learning rate schedule!  │ exact failing layer node!│
└──────────────────────────┘

Step 1: Check Input Data for Unhandled NaNs or Outliers

Before training, assert that input feature matrices contain zero NaNs or Infinities:

assert not torch.isnan(inputs).any(), 'NaNs detected in input features!'
assert not torch.isinf(inputs).any(), 'Infinities detected in input features!'

Confirm that numerical features are properly scaled. Raw values like $10,000,000$ cause numeric instability during matrix multiplication.

Step 2: Implement Gradient Clipping

Unstable loss surfaces cause Exploding Gradients. Cap parameter gradient norms before invoking optimizer.step():

loss.backward()
# Clip gradient norms to a maximum threshold of 1.0
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()

Step 3: Add Epsilon Constants to Logarithms and Square Roots

Computing $\log(0)$ or $\frac{1}{\sqrt{0}}$ produces negative infinity or NaN values:

# BAD: Crashes if probs contains 0.0!
loss = -torch.log(probs)

# GOOD: Add tiny epsilon offset or use torch.log(probs.clamp(min=1e-12))
loss = -torch.log(probs + 1e-12)

Use combined loss functions like nn.CrossEntropyLoss() or nn.BCEWithLogitsLoss() which integrate numerical stabilization internally!

Step 4: Lower the Learning Rate and Use Warm-Up

A learning rate that is too high causes parameters to overshoot loss valleys into unstable steep regions. Reduce learning rate by $10\times$ and use a Learning Rate Warm-Up Schedule.

Step 5: Enable PyTorch Autograd Anomaly Detection

To pinpoint the exact forward or backward layer operation that introduced NaN:

# Enable anomaly detection (Use during debugging only, slows down training!)
torch.autograd.set_detect_anomaly(True)

PyTorch will print a detailed stack trace highlighting the exact layer operation that output NaN!

Say this out loud

Debugging NaN loss requires checking input feature data for unhandled nulls, clipping exploding gradient norms, adding epsilon offsets to logarithms, and lowering high learning rates. Enabling PyTorch autograd anomaly detection pinpoints the exact layer operation that introduced numerical instability.

Followups to expect

1. Why does FP16 Mixed Precision training cause NaN loss more frequently than FP32? FP16 has a narrow dynamic range ($6 \times 10^{-5}$ to $65504$), causing small gradients to underflow to zero and large gradients to overflow to infinity. Use GradScaler or BF16 precision to resolve this.
2. What is weight initialization impact on NaN loss? Initializing weights with values that are too large causes activations to explode exponentially across deep layers, leading to instant NaN loss on epoch 1.