Neural networks, backpropagation and numerical stability

Why it matters

Frameworks compute gradients for you, but bugs and instabilities live in what they hide: a softmax that overflows, a loss that turns into NaN, gradients that vanish in a deep stack, a batch-norm layer left in training mode at inference. Typical probes in ML engineering and research interviews are a backprop derivation on paper or a "why does this loss become NaN?" question.

Core ideas

Layers and nonlinearity. An MLP alternates affine maps h = W x + b with nonlinearities such as ReLU or GELU. Without the nonlinearities, a stack of linear layers collapses into a single linear map.

Backpropagation. Treat the network as a graph of primitive operations. The forward pass computes outputs and stores what the backward pass will need. The backward pass applies the chain rule from the loss backwards: each node receives the gradient of the loss with respect to its output and multiplies it by its local derivative, a vector-Jacobian product, without ever building a full Jacobian. For Y = X W with X of shape (n, d) and W of shape (d, k): dL/dW = X^T dL/dY and dL/dX = dL/dY W^T; matching shapes tells you where the transposes go. When a value feeds several consumers, its incoming gradients add. Reverse mode costs a small constant multiple of one forward pass however many parameters there are, which is why it suits a scalar loss over millions of weights.

Autograd in practice. PyTorch records operations on tensors that require gradients and saves the tensors each backward formula needs, which is why stored activations are often the largest part of training memory. Modifying a saved tensor in place makes backward fail with an error; torch.no_grad() and inference mode skip graph recording. Gradients accumulate into .grad across backward calls until you zero them.

Gradient checking. Compare analytic gradients with central differences (f(x + h) - f(x - h)) / (2h) in float64, using relative error. It is the fastest way to validate a hand-written backward pass.

Numerical stability. In float32, exp overflows for inputs above roughly 88. Softmax is unchanged by subtracting a constant from every logit, so compute it as exp(z - max(z)) / sum_j exp(z_j - max(z)). For log-probabilities use log_softmax(z) = z - logsumexp(z), where logsumexp(z) = m + log(sum_j exp(z_j - m)) and m = max(z). Prefer the framework's fused cross-entropy on raw logits to taking the log of a softmax output, which can produce log(0); likewise use the with-logits form of binary cross-entropy.

Initialization. Weights that are too large make activations and gradients grow layer by layer; too small, they shrink towards zero. Scaling the weight variance inversely with the number of inputs (or with the average of inputs and outputs) keeps signal variance roughly constant: Xavier/Glorot (Var(w) = 2 / (fan_in + fan_out)) for tanh-like units, He/Kaiming (Var(w) = 2 / fan_in) for ReLU, which zeroes about half its inputs. Initializing every weight in a layer to the same value keeps the units identical forever, so they never learn different features.

Normalization. Batch norm normalizes each feature across the batch (and spatial positions), keeps running statistics for inference, and therefore behaves differently in train and eval mode; it struggles with small or variable batches and with sequences. Layer norm normalizes across the features of each example, so it is independent of batch size and identical at training and inference; transformers use it or RMSNorm, which skips mean-centering. Both learn a scale and usually a shift.

Vanishing and exploding gradients. Backpropagating through many layers multiplies many Jacobians; if their norms sit below or above one, gradients shrink or grow exponentially with depth. Remedies: ReLU-family activations, careful initialization, normalization, residual connections (which give gradients an identity path), and gradient clipping against explosions.

Common misconceptions

• "Backprop computes the full Jacobian of every layer." It propagates vector-Jacobian products; materializing Jacobians would be far too expensive.
• "Subtracting the max changes the softmax." It does not, because softmax is shift-invariant; it only prevents overflow.
• "Batch norm and layer norm are interchangeable." They normalize over different axes and differ in train and eval behavior; batch norm with tiny batches, or evaluated with stale running statistics, is a classic bug.
• "Initializing all weights to zero is a safe start." It keeps hidden units identical; biases may start at zero, weights may not all be equal.
• "A NaN loss means the learning rate is too high." Often, but also check log(0), division by a zero variance, float16 overflow and corrupt inputs.

How to approach it

Derive a two-layer MLP backward pass on paper, implement it in NumPy and gradient-check it. Build a tiny autograd engine to understand graph recording and gradient accumulation. Read the framework notes on what is saved for backward. In interviews, write shapes next to every equation and present stability tricks as consequences of floating-point limits, not as rules to memorize.