ML Debugging — Deep Dive

Frontier-lab interview prep. Pair with INTERVIEW_GRILL.md.

Senior ML interviews increasingly include “the model isn’t training, debug it” — testing whether you have a systematic methodology, not just textbook knowledge. This deep dive is the methodology: how to debug failing training, NaN gradients, leakage, and production regressions.

1. The debugging tree

When something is wrong, work through layers in order. Most failures hit at one of these:

1. Data
2. Pipeline (preprocessing, batching, augmentation)
3. Model (architecture, init, forward pass)
4. Loss (function, label format, reduction)
5. Optimizer (LR, momentum, weight decay)
6. Training loop (gradient clip, accumulation, AMP)
7. Evaluation (metric, split, leakage)
8. Deployment (serving, infra)

The principle: cheap checks first. Plot losses, eyeball data, sanity-check shapes. Don’t dive into custom gradient computations until you’ve ruled out trivial bugs.

Saying it out loud. My rule is: debug in the order that things are cheap to check, not in the order they’re interesting. Data first, then the pipeline, then the model, then loss, optimizer, training loop, evaluation, and only then infrastructure. The reason is base rates — the overwhelming majority of “the model won’t train” turns out to be a label format, a shape, or a learning rate, and almost none of it is an exotic gradient bug. So I plot the loss, print a batch, and check shapes before I ever open the modelling code. The failure mode I’m avoiding is the classic one: two days rewriting an attention implementation when the labels were off by one.

2. Loss curve interpretation

“Loss going down — am I done?”

No. Sanity-check:

• Train loss should decrease.
• Val loss should track or lag.
• Eventually val loss flattens or rises (overfitting).

Common loss-curve patterns

Loss not decreasing:

• Bad LR (too low → no movement; too high → bouncing without progress).
• Frozen weights (forgot requires_grad=True).
• Disconnected gradients (broken computation graph).
• Wrong loss function or label format.
• Too-small dataset (model trivially fits but val won’t improve).

Loss exploding (NaN):

• LR too high.
• Numerical overflow (FP16 issue).
• Bad weight init.
• Instability in attention softmax (large logits).
• Division by zero somewhere.

Loss decreasing but val not improving:

• Severe overfit. More data, regularize, smaller model.
• Train-val mismatch (preprocessing, distribution).
• Leakage in train but not val (different shape of leakage).

Loss flatlines high then jumps:

• “Phase transition” — sometimes models break through plateau. Common in RL.
• Or warmup not finished.

Loss decreasing then suddenly spikes:

• Bad batch (OOD example, very long sequence, gradient cliff).
• Optimizer state went bad.
• Gradient clipping not applied.
• Re-warmup might be needed.

Quick debugging actions

• Plot with log y-axis to see early dynamics.
• Compare loss curves across runs (changed one thing — should affect curve in expected way).
• Look at per-batch loss, not just smoothed running average.

Saying it out loud. The loss curve is the cheapest diagnostic you own, and each shape means something specific. Dead flat from step zero says the gradient isn’t reaching the weights — frozen parameters, a detached graph, or a learning rate so small nothing moves. Noisy and refusing to descend usually means the learning rate is too high. Exploding to NaN means overflow or a step off a cliff. Train falling while validation climbs is plain overfitting, and both being suspiciously good is a leakage smell. Two habits that pay: plot on a log y-axis so you can see the first hundred steps, and look at per-batch loss rather than the smoothed average, because a single bad batch spikes the raw curve and vanishes from the moving average.

3. Data-side debugging

“Loss looks weird, model isn’t learning”

Sanity checks:

• Does the data look right? Plot a few examples.
• Are labels in expected range / format?
• Are images normalized properly? (Common bug: $/255$ vs Standardization with mean/std.)
• Do tokens decode back to original text?

Sanity check 1: overfit a single batch

Take one batch. Train on it for many steps. Loss should go to ~0.

If it can’t even fit one batch:

• Wrong loss function or label format.
• Frozen layer.
• Too small a model.
• Bug in data loading.

This 5-minute test catches many issues immediately.

Sanity check 2: tiny dataset

Take 100 examples. Train. Should overfit quickly (train accuracy ~100%).

If it doesn’t, model lacks capacity or there’s a fundamental bug.

Sanity check 3: data inspection

Print 5 random batches. Look at shapes, label distributions, raw values.

If something looks off, fix that first before assuming a model issue.

Saying it out loud. Before I blame the model I try to fit one batch. Take a single batch, turn off regularization and shuffling, and train on it for a few hundred steps — the loss should crater toward zero, because a model that can’t memorize thirty-two examples has a bug, not a capacity problem. If that passes, I scale up to a hundred examples and expect near-perfect training accuracy. Then I print raw batches and look at them: shapes, value ranges, label distribution, and for text, decode the tokens back and read them. That whole loop is about five minutes and it catches most real bugs. The classic one it finds is a normalization mismatch — dividing by 255 in training and standardizing with ImageNet mean and standard deviation at inference.

4. NaN debugging

Causes

• FP16 overflow: $e^x$ for $x>11$ overflows in FP16 (max finite value 65504); $x>88$ is the FP32 threshold. Use BF16 (extended exponent range) or stable softmax.
• Division by zero: variance estimate hits 0; output of LayerNorm; division in normalization.
• Log of zero: $\log p$ for $p=0$. Add $ϵ$ (e.g., $\log (p+10^{-9})$).
• Square root of negative: numerical drift makes a “non-negative” value slightly negative.
• Inf gradient: explodes through layers due to bad init or large input.

Detecting NaN early

• Loss is NaN → too late, weights already corrupted.
• Add assert not torch.isnan(x).any() after suspect operations.
• PyTorch has torch.autograd.set_detect_anomaly(True) — slow but catches first NaN site.

Triage

• When did it appear? Step 0? Step 5000?
• Step 0: bad init or first batch issue.
• Later: optimizer instability, bad batch, explosion.

Fix patterns

• Gradient clipping (norm 1.0).
• Lower LR.
• BF16 over FP16.
• Compute attention in higher precision.
• Add $ϵ$ in normalizations.
• Restart from earlier checkpoint.

Saying it out loud. With a NaN, the first question is always when — step zero or step five thousand — because they mean completely different things. At step zero it’s initialization, a bad first batch, or an input that’s already broken. Thousands of steps in, it’s an instability that finally got triggered: a huge attention logit, a divide by a variance that hit zero, a log of zero, or an outlier batch producing a giant gradient. By the time the loss reads NaN the weights are already poisoned, so you catch it earlier with assertions after suspect ops or with anomaly detection, which is slow but points at the first bad operation. The fix list in order: clip gradients at norm 1.0, lower the learning rate, switch FP16 to BF16 — same memory, far more exponent range, no loss scaling needed — and restart from the last clean checkpoint.

5. Leakage debugging

Symptoms

• Offline metrics suspiciously high.
• Train + val accuracy both 99%, test in production poor.
• Model “feature importance” shows a feature that shouldn’t exist.

Classic forms

• Target leakage: feature computed using target (or post-target).
• Train-test contamination: same record in both splits.
• Preprocessing leakage: stats computed on full data.
• Group leakage: same user/patient on both sides.
• Temporal leakage: future used to predict past.

Detection

Suspicious AUC:

• Train classifier with features sorted by importance; if top-1 feature alone gives AUC > 0.95, it’s probably leakage.
• Train without each feature individually; one with huge drop → suspect.

Cross-correlation check:

• Compute correlation between every feature and the label.
• If correlation > 0.9 for any feature, audit it.

Held-out time-period validation:

• If you have time-stamped data, hold out the last $X\%$ as validation.
• If accuracy drops a lot relative to random split, you had temporal leakage.

Common bug example

# WRONG: scaler fit on full dataset before split
scaler = StandardScaler().fit(X)  # uses test statistics
X = scaler.transform(X)
X_train, X_test = train_test_split(X, ...)

# RIGHT: split first, fit on train only
X_train, X_test = train_test_split(X, ...)
scaler = StandardScaler().fit(X_train)
X_train = scaler.transform(X_train)
X_test = scaler.transform(X_test)

Saying it out loud. Leakage is when information the model wouldn’t have at prediction time sneaks into training, and the tell is a number that’s too good. If a single feature gets you AUC above 0.95, or any feature correlates above 0.9 with the label, I go audit it before I celebrate. The forms worth naming out loud are target leakage, where a feature is computed downstream of the outcome; group leakage, where the same user or patient lands in both splits; temporal leakage, where you use the future to predict the past; and preprocessing leakage, the most common by far — fitting a scaler or an imputer on the whole dataset before splitting, so test statistics bleed into training. The diagnostic that catches temporal leakage is comparing a random split against a time-ordered holdout: if accuracy falls off a cliff on the time split, you had it.

6. Gradient checking

For custom layers / losses, verify gradients numerically.

def gradient_check(f, x, eps=1e-5):
    """Compare analytical gradient to numerical."""
    analytical = f.backward(x)
    numerical = np.zeros_like(x)
    for i in np.ndindex(x.shape):
        x_plus = x.copy(); x_plus[i] += eps
        x_minus = x.copy(); x_minus[i] -= eps
        numerical[i] = (f.forward(x_plus) - f.forward(x_minus)) / (2 * eps)
    rel_error = np.abs(analytical - numerical) / (np.abs(analytical) + np.abs(numerical) + eps)
    return rel_error.max()

If rel_error > $10^{-5}$, suspect a bug.

PyTorch has torch.autograd.gradcheck(func, inputs). Use it for custom autograd functions.

Saying it out loud. Gradient checking is how you prove a hand-written backward pass is right: perturb one input by a tiny epsilon in each direction, take the two-sided finite difference of the forward pass, and compare against the analytical gradient. You compare relative error, not absolute, and anything under about $10^{-5}$ is fine while anything above $10^{-3}$ means a real bug. Two practical notes: use double precision, because in float32 the noise floor swamps the signal, and avoid checking at kinks like ReLU at zero, where the finite difference is legitimately meaningless. It’s $O(n)$ forward passes, so you run it once on a tiny tensor, not in training — in PyTorch, torch.autograd.gradcheck does exactly this.

7. Distribution-shift debugging

Symptoms

• Offline metrics good, online metrics bad.
• Model degrades over time.
• Subgroup performance worse than aggregate.

Investigation

• Compare input distributions: train vs production. KS test, KL, PSI per feature.
• Compare prediction distributions: train vs production. Output histograms shifted?
• Compare actual labels (where available): production positive rate vs train.
• Subgroup analysis: performance by user segment / region / device.

Common causes

• Covariate shift: new user demographics. Reweight or retrain.
• Concept drift: relationship $p(y|x)$ evolves. Retrain on fresh data.
• Selection bias: only certain populations seen offline.
• Pipeline drift: feature definitions changed silently.

Mitigation

• Online retraining cadence.
• Feature monitoring with alerts.
• Shadow / canary deployment to catch regressions early.

Saying it out loud. When offline looks great and production looks bad, my first hypothesis is that the data moved, not that the model broke. I compare three distributions between training and live traffic: the inputs feature by feature, the model’s own output scores, and the label rate wherever labels arrive. A shifted score histogram is the fastest signal because it’s one plot and needs no labels. Then I name which kind of shift it is: covariate shift means the inputs moved but the input-to-label relationship held, so reweighting or retraining fixes it; concept drift means the relationship itself changed and only fresh labels help. The one that bites hardest is pipeline drift — an upstream team silently redefines a feature and the model degrades with no distribution alarm at all, which is why you monitor features and ship behind a canary.

8. Common interview gotchas

Saying it out loud. The thread through this table is that in a debugging round, interviewers score your ordering, not your knowledge of any single fix. Reaching for a bigger model or a retrain first reads as guessing; cheap diagnostics first reads as experience. So for a flat loss I overfit one batch, for a suspiciously good AUC I suspect leakage, for a production regression I roll back first and investigate second, and for a NaN I ask when it appeared. Say the diagnostic before the fix every time. The one concrete number worth having ready: gradient clipping at global norm 1.0 is the standard for transformers.

9. Eight most-asked debugging questions

1. Loss is not decreasing — walk through your investigation. (Overfit one batch; verify loss; check LR; inspect data.)
2. NaN gradient — what’s your debugging process? (When did it appear; FP16/overflow; inf in attention; gradient clip; restart from checkpoint.)
3. Train accuracy 99%, test 60% — what’s wrong? (Severe overfit. Or leakage in train. Or distribution mismatch.)
4. Offline AUC 0.95, online accuracy bad — what to check? (Position bias, distribution shift, label time leakage, counterfactual issue.)
5. Model regressed in production — investigation. (Rollback first. Then: data, features, infra, drift.)
6. Implement gradient checking for a custom layer. (Numerical: $(f(x+ϵ)-f(x-ϵ))/(2ϵ)$.)
7. What does a flat loss curve mean? (LR issue; warmup not finished; frozen weights; phase transition pending.)
8. What sanity checks before training? (Overfit one batch; tiny dataset to 99%; inspect data.)

Saying it out loud. If I get any of these, I answer with a process rather than a guess, and I say the process out loud in order: form a hypothesis, name the cheapest test that could disprove it, run it, then narrow. Loss not moving? Overfit one batch. NaN? Ask when it first appeared. Ninety-nine train against sixty test? Decide between overfitting, leakage, and a train-serving distribution mismatch, and name the check that distinguishes them. Offline good but online bad? Distribution shift and feedback effects like position bias before anything else. The tradeoff I state explicitly is speed versus certainty in production: roll back first to stop the bleeding, then debug on the artifacts, because a live regression is not the place to be curious.

10. Drill plan

• For each common loss-curve pattern (flat, exploding, val-train gap), recite cause + diagnostic + fix.
• For NaN: list 5 causes and the corresponding fix.
• For leakage: code the “fit scaler before split” bug and the corrected version.
• Practice the “overfit one batch” sanity check on a real model.
• Time yourself: 5 min to outline a debugging investigation given a vague problem.

11. Further reading

• Karpathy, A Recipe for Training Neural Networks (2019 blog) — the canonical practical guide.
• Goodfellow, Bengio, Courville, Deep Learning, ch. 11 — practical methodology.
• Smith (2018), A disciplined approach to neural network hyperparameters.
• Andrej Karpathy’s tweets and lectures on debugging (timeless).
• Hidden Technical Debt in ML Systems (Sculley et al. 2015) — production-side issues.