Training deep networks

The practical layer: what to set, what breaks, and how to diagnose it. Most interview questions here are debugging scenarios.

Initialisation

Weights initialised badly make a network untrainable regardless of everything else. The goal is keeping activation variance roughly constant across layers — too small and the signal vanishes, too large and it explodes.

He initialisation accounts for ReLU killing half the activations, hence the factor of 2 on fan_in alone. Frameworks default sensibly, so this matters mainly when writing custom layers.

Never initialise all weights to zero — every neuron in a layer computes the same thing and receives the same gradient, so the layer collapses to a single unit forever. Biases at zero are fine.

Learning rate: the parameter that matters most

Find a starting point with an LR range test — increase the LR exponentially over a few hundred steps and plot loss against LR. Pick roughly an order of magnitude below where it starts diverging.

Typical values: 3e-4 for transformers with AdamW, 1e-3 for small networks, 1e-5 to 5e-5 for fine-tuning a pretrained model.

Schedules and warmup

from torch.optim.lr_scheduler import OneCycleLR

scheduler = OneCycleLR(optimizer, max_lr=3e-4,
                       total_steps=total_steps, pct_start=0.05)

Linear warmup then cosine decay is the standard transformer recipe. Warmup exists because early gradients are large and Adam’s second-moment estimate isn’t yet reliable — stepping at full rate immediately can destabilise training permanently. Decay exists because you want large steps to find the basin and small steps to settle in it.

Schedules are defined in steps, not epochs. Changing the batch size changes the number of steps, so the schedule silently changes with it.

Gradient problems

Exploding — loss spikes or goes NaN. Clip:

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

Standard in transformer and RNN training. Clip by global norm, not per-parameter.

Vanishing — early layers stop learning. Fixes are architectural rather than tuning: ReLU-family activations, residual connections (which give the gradient an identity path straight through), and normalisation layers.

Diagnose by logging gradient norms per layer. A gradient norm several orders of magnitude smaller at layer 1 than at layer 20 is the signature.

for name, p in model.named_parameters():
    if p.grad is not None:
        print(name, p.grad.norm().item())

Mixed precision

Standard practice, and worth being able to explain:

from torch.amp import autocast, GradScaler

scaler = GradScaler()
for xb, yb in loader:
    optimizer.zero_grad()
    with autocast(device_type="cuda", dtype=torch.bfloat16):
        loss = criterion(model(xb), yb)
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    scaler.step(optimizer)
    scaler.update()

Roughly halves memory and speeds up matmuls substantially on modern GPUs.

fp16 vs bf16 is the question worth knowing. fp16 has more mantissa bits but a narrow exponent range, so small gradients underflow to zero — which is why GradScaler exists, multiplying the loss to keep gradients representable. bf16 has the same exponent range as fp32, so it doesn’t underflow and doesn’t strictly need loss scaling. On hardware that supports it, bf16 is the safer default; fp16 remains common on older GPUs.

Note unscale_ before clipping — clipping scaled gradients clips the wrong magnitude.

Memory

Training memory is roughly:

weights + gradients + optimizer state (2x for Adam) + activations

The first four scale with parameters; activations scale with batch size and sequence length. Levers, in order of what to try:

Gradient checkpointing is the classic compute-for-memory trade: discard intermediate activations and recompute them during the backward pass.

Reproducibility

torch.manual_seed(42)
np.random.seed(42)
torch.use_deterministic_algorithms(True)     # slower; some ops unsupported

Full determinism on GPU costs performance and isn’t always achievable — some cuDNN kernels are non-deterministic by design. Seed everything, pin library versions, and expect small run-to-run variation. Report a mean over several seeds rather than a single number when comparing architectures; the difference between two runs of the same config is often larger than the difference between two configs.

A debugging order that works

1. Overfit a single batch. If the model can’t drive loss to ~0 on 8 examples, the bug is in the model, loss or data pipeline — not the hyperparameters. This is the highest-value first test and it takes a minute.
2. Check shapes and the loss input. Logits vs probabilities, label dtype, class indices vs one-hot.
3. Check the data. Visualise a batch after augmentation. Confirm labels align with inputs.
4. Drop the learning rate 10x. Resolves a surprising share of instability.
5. Turn off regularisation and augmentation, confirm it trains, then reintroduce.
6. Check normalisation. Inputs standardised; model.eval() used at validation.

Interview angle

• “Your loss goes to NaN. Walk me through debugging.” — learning rate first (drop 10x), then check for log(0) or division by near-zero in the loss, then gradient clipping, then whether fp16 underflow is involved (try bf16), then input normalisation and bad data rows.
• “Why warmup?” — early gradients are large and Adam’s variance estimate is uncalibrated, so full-rate steps at the start can destabilise training irrecoverably. Warmup ramps in gently.
• “fp16 or bf16?” — bf16 where supported: it has fp32’s exponent range so gradients don’t underflow and loss scaling isn’t strictly needed. fp16 has better precision but a narrow range, which is exactly why GradScaler exists.
• “How do you train a model that doesn’t fit in GPU memory?” — gradient accumulation with a smaller batch, mixed precision, gradient checkpointing, an 8-bit optimiser, LoRA instead of full fine-tuning, and FSDP/ZeRO sharding across devices. In roughly that order of effort.
• “Why not initialise weights to zero?” — every unit in a layer would compute identically and receive identical gradients, so the layer never differentiates. Symmetry has to be broken by random initialisation.
• “First thing you do when a model won’t learn?” — try to overfit a single small batch. Success narrows the problem to optimisation or regularisation; failure points at the model, loss or data pipeline.