Training Deep Networks: Initialization, Normalization & Optimizers

Training a deep network means steering millions of coupled parameters through a non-convex loss surface using noisy mini-batch gradients. Most failures are not explained by “the optimizer did not find the global minimum.” They arise because signals vanish or explode, learning rates mismatch scale, batches are corrupted, loss/labels disagree, or the model fits shortcuts.

Signal scale across layers #

For a layer $z_j=\sum_{i=1}^{n}w_{ji}x_i$, independent zero-mean terms give roughly

$$\mathop{\mathrm{Var}}(z_j)\approx n\mathop{\mathrm{Var}}(w)\mathop{\mathrm{Var}}(x).$$

If $n\mathop{\mathrm{Var}}(w)$ is much larger than one, activations grow across layers; if smaller, they shrink. Gradients multiply analogous Jacobians in reverse, causing explosion or vanishing.

Glorot/Xavier initialization balances forward and backward variance for roughly symmetric activations, often using variance $2/(fan_{in}+fan_{out})$. For ReLU, about half of a symmetric preactivation is zeroed, motivating He initialization with variance about $2/fan_{in}$. These are model-dependent starting scales, not guarantees.

For a linear layer with 100 independent unit-variance inputs, weights with variance 1 would produce preactivation variance near 100, so repeated layers would rapidly explode. Weights with variance $1/100$ keep linear preactivation variance near one; a ReLU layer motivates roughly $2/100$ because rectification discards about half the second moment. Correlation and learned nonzero means weaken this calculation, which is why activation statistics still need measurement during training.

Residual connections provide short identity paths, letting layers learn corrections rather than a complete transformation and giving gradients routes that avoid long products of poorly conditioned Jacobians.

Activation functions change gradient flow #

Sigmoid derivatives become tiny in saturated tails; deep products then vanish. Tanh is zero-centered but still saturates. ReLU has derivative one on its positive side, improving flow, but a unit kept negative for all inputs receives zero gradient. Leaky/smooth variants alter that trade-off.

Monitor activation distributions and the fraction of zero/saturated units. Replacing an activation without retuning initialization and learning rate changes the entire scale system.

Normalization is architecture-dependent #

Batch normalization standardizes channel activations using mini-batch statistics during training, then uses running estimates for inference. It often permits larger learning rates and smooths optimization behavior. The original “internal covariate shift” explanation is not the only accepted account of its benefit; treat the observed optimization effect as stronger evidence than that single causal story.

Small, non-IID, or sequence-correlated batches make batch statistics noisy and create train/inference mismatch. Layer normalization normalizes within each example and is independent of batch composition, which suits sequence models. Normalization also changes what parameter scaling means and interacts with weight decay.

Always verify evaluation mode. Serving with training-mode batch normalization makes predictions depend on other requests in the batch.

SGD, momentum, and Adam #

Mini-batch SGD updates $\theta_{t+1}=\theta_t-\eta g_t$. Momentum accumulates a velocity, damping oscillation in steep directions and continuing along persistent gradients.

Adam tracks exponential moving averages of gradients and squared gradients:

$$m_t=\beta_1m_{t-1}+(1-\beta_1)g_t,$$

$$v_t=\beta_2v_{t-1}+(1-\beta_2)g_t^2,$$

then applies bias-corrected $\hat m_t/(\sqrt{\hat v_t}+\epsilon)$. It adapts step sizes per parameter and often reaches a useful regime quickly. SGD with momentum can generalize competitively and may be easier to reason about in some vision regimes; optimizer choice is empirical, not hierarchical.

For adaptive optimizers, AdamW decouples weight decay from the loss-gradient update. Adding an $L_2$ gradient and multiplying weights by a decay factor are equivalent under plain SGD scaling but not under Adam's per-coordinate normalization.

Learning rate is the central control #

Too large: loss diverges, oscillates, or produces non-finite values. Too small: progress is glacial and optimization may stay in plateaus. Warmup prevents early unstable updates before statistics and adaptive moments settle; decay schedules reduce step size for refinement. A learning-rate range test can reveal a useful scale, but the final schedule still needs validation.

Gradient clipping limits a global norm or individual values and is especially useful for recurrent/unstable sequences. It contains an explosion; it does not explain or remove its cause.

A disciplined training loop #

1. Overfit a tiny batch to verify model, loss, labels, and gradient wiring.
2. Compare against a simple baseline and shuffled-label control.
3. Log training/validation loss, task metrics, learning rate, gradient norms, activation statistics, and throughput.
4. Use checkpointed validation for early stopping; keep the final test untouched.
5. Record random seeds, data snapshot, preprocessing, code, hardware-relevant determinism settings, and model state.

Exact bitwise reproducibility can be costly or unavailable across hardware kernels. The requirement is to identify nondeterminism, measure run-to-run variance, and preserve enough provenance to reproduce the statistical conclusion.

Regularization and data #

Weight decay limits parameter scale; dropout samples thinned subnetworks during training; augmentation encodes label-preserving invariances; early stopping limits how far fitting proceeds. Their effects overlap. Applying all defaults can underfit or distort uncertainty.

Train/validation curves diagnose whether more representative data, stronger regularization, or more capacity is plausible. A falling training loss with flat validation loss is not an optimizer failure; continuing longer usually strengthens the wrong fit.

Failure triage #

• Non-finite loss: inspect inputs, target ranges, logits/loss pairing, learning rate, mixed-precision scaling, and gradient norms.
• No learning: verify labels, nonzero gradients, trainable parameters, activation saturation, and baseline loss.
• Training learns, validation fails: audit split/leakage, augmentation, regularization, and population mismatch.
• Evaluation changes with batch size/order: inspect batch normalization, padding masks, stateful recurrence, and stochastic layers in eval mode.

Optimization is an observable system. Treat “try another optimizer” as a hypothesis after gradients, scales, data, and modes have been measured.