9.3 Initialization, Optimization, and the Training Loop: Loss Reduction Is Just the First Vital Sign of a Training System

The computation graph can provide gradients, but the model workshop remains unstable: a slightly high learning rate leads to NaN values, swapping the seed produces different outcomes, and validation performance degrades even as training loss continues to drop.

Training is a system: initialization sets the initial signal and gradient scale, the optimizer transforms gradients into parameter updates, the data loader introduces noise and determines throughput, mode controls dropout and normalization, and validation dictates when to save checkpoints and when to stop.

Learning Objectives

• Choose initialization strategies based on activation function, fan-in, and fan-out;
• Understand the differences in update rules between SGD, Momentum, Adam, and AdamW;
• Correctly manage learning rate, batch size, normalization, and gradient clipping;
• Write a minimal and reliable PyTorch training loop that includes train and eval modes;
• Record checkpoints, maintain reproducibility, and track diagnostic signals.

1. Why Zero-Weight Initialization Fails

If hidden units in the same layer start with identical weights and receive the same gradient, they will remain indistinguishable and fail to develop distinct features. This is the symmetry problem.

Initializing biases to zero typically avoids this issue, because random weights break the symmetry among hidden units; the claim that "no parameter should be zero" is not strictly valid.

Random initial values also need to control variance. If the variance of activation or gradient signals grows or shrinks across layers, it can lead to saturation, explosion, or vanishing gradients.

2. Xavier and Kaiming's Assumptions

Xavier/Glorot initialization adjusts variance based on fan-in and fan-out, and is commonly paired with approximately symmetric activations like tanh. Kaiming/He initialization accounts for the fact that ReLU-like activations only respond to a subset of inputs, preserving forward variance based on fan-in.

These are approximations derived under assumptions of independent, zero-mean inputs and specific activation distributions, neither universal nor architecture-agnostic theorems. Residual branches, attention mechanisms, gating units, normalization layers, and very deep networks may require specialized scaling.

PyTorch modules each have their own default initialization schemes. Do not assume that a module is "He-initialized" without verifying the current version and actual implementation. If custom initialization is applied:

from torch import nn

def initialize(module):
    if isinstance(module, nn.Linear):
        nn.init.kaiming_normal_(module.weight, nonlinearity="relu")
        if module.bias is not None:
            nn.init.zeros_(module.bias)

model.apply(initialize)

If the final layer or activation function differs, do not indiscriminately apply the same initialization scheme.

3. SGD and Momentum

The mini-batch gradient $g_t$ is a noisy estimate of the full gradient under the corresponding sampling condition. Basic stochastic gradient descent (SGD):

$$ \theta_{t+1} = \theta_t - \eta g_t. $$

Momentum accumulates an exponentially weighted direction (specific symbol or implementation conventions vary):

$$ v_t = \mu v_{t-1} + g_t, \quad \theta_{t+1} = \theta_t - \eta v_t. $$

This can accelerate progress in a consistent direction and smooth out oscillations in alternating directions, but it does not guarantee avoidance of saddle points or improved generalization. The trajectory is jointly determined by the learning rate, momentum, batch size, and schedule.

4. Adam and AdamW

Adam maintains exponentially weighted estimates of the first-moment (gradient mean) and second-moment (gradient squared) statistics, then applies bias correction and scales the update by coordinate. It often converges quickly to usable training results, especially when dealing with sparse or heterogeneous gradient magnitudes, though it is not universally the optimal optimizer for all tasks.

In Adam, incorporating an L2 penalty directly into the gradient does not fully equate to decoupled weight decay. AdamW separates weight decay from the adaptive gradient update process. Additionally, decisions about which parameters to apply decay to (such as biases or normalization scale parameters) are often made based on architectural or experimental design, with these parameters typically excluded from decay (no-decay group). This choice remains a design decision specific to the model and experiment.

When comparing optimizers, it's essential to apply comparable hyperparameter budgets, learning rate schedules, and computational resources, rather than simply comparing AdamW against SGD with its default learning rate.

5. Learning Rate is one of the primary stability parameters

Too high: loss oscillates or diverges, activation or gradient values become Inf or NaN; too low: the optimizer makes minimal progress within a finite budget and may get stuck in a suboptimal region.

Common schedules include warmup, step or exponential decay, cosine decay, and plateau-driven decay. Warmup can reduce update steps during large batch training, adaptive optimization, or early stages of instability, but it is not required for all small networks.

"Using Adam with 0.001" is merely a starting point for certain configurations. The effective learning rate is also influenced by loss reduction, batch size, gradient accumulation, parameterization, and numerical precision.

6. Changing Batch Size and System Behavior

Small batch: higher gradient noise, more frequent updates, potentially lower throughput. Large batch: may utilize hardware more efficiently, fewer updates per epoch, higher memory usage, and might require adjustments to learning rate or scheduling.

"The larger batch size always leads to worse generalization" is not a law. When comparing, total examples, number of updates, training schedule, and compute resources must be controlled. Gradient accumulation simulates a larger effective batch size for gradient averaging, but it cannot fully replicate BatchNorm statistics, optimizer update frequency, or the sequence of random augmentations.

7. Normalization Goes Beyond Preventing Gradient Explosion

BatchNorm

During training, mini-batch statistics are used to normalize the data, and running estimates are updated in real time. During evaluation, the precomputed running statistics are used instead. Results can be significantly affected by small or non-iid batches, distributed data sharding, and transitions between training and evaluation modes.

LayerNorm / RMSNorm

These layers normalize along the feature dimension of a single sample, independent of batch statistics, making them well-suited for sequence models. However, the exact normalized shape and data type (dtype) must still be carefully verified for each implementation.

The placement of normalization layers (before or after activation functions, or before/after residual branches) has no universal solution across all architectures. Pre-norm and post-norm configurations alter the optimization path. Always follow the design of the target architecture and validate choices through ablation studies. Avoid blindly adhering to the "Linear–BN–ReLU" pattern as a one-size-fits-all approach.

8. Different Roles of Regularization

• Weight decay: Shrinks parameter values; its effect depends on the parameterization and optimizer used.
• Dropout: Randomly zeros out neurons during training (with scaling), and is disabled during evaluation.
• Data augmentation: Introduces assumptions of task invariance by adding diverse training samples.
• Label smoothing: Modifies the target distribution, improving calibration and reducing overconfidence.
• Early stopping: Selects the optimal number of training epochs based on validation performance.
• Stochastic depth and mixup, among others: Alter the objective function or model architecture in distinct ways.

These are not interchangeable "overfitting prevention add-ons" that can be freely stacked. First, identify the specific failure mode, then tune the strength of each technique within the bounds of validation performance.

9. Gradient Clipping Is a Fence, Not a Repairman

Global norm clipping:

$$ g\leftarrow g\cdot\min\left(1,\frac{c}{\lVert g\rVert}\right). $$

It limits the norm of the gradient per step and is commonly used in sequence models or unstable training scenarios. If the clipping threshold is too small, it may continuously alter the optimization direction or scale. If clipping is frequently triggered, investigate potential issues such as loss scaling, data outliers, initialization, learning rate, or numerical overflow.

Record the gradient norm before clipping and the proportion of times clipping is triggered. Clip only before optimizer.step(); when using mixed precision, first unscale according to the current AMP API, then apply clipping.

10. Minimal Reliable Training Loop

import torch
from torch import nn

def train_one_epoch(model, loader, optimizer, loss_fn, device):
    model.train()
    total_loss = 0.0
    total_examples = 0

    for features, target in loader:
        features = features.to(device)
        target = target.to(device)

        optimizer.zero_grad(set_to_none=True)
        logits = model(features)
        loss = loss_fn(logits, target)

        if not torch.isfinite(loss):
            raise FloatingPointError(f"non-finite loss: {loss.item()}")

        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0)
        optimizer.step()

        batch = target.shape[0]
        total_loss += loss.item() * batch
        total_examples += batch

    return total_loss / total_examples

def evaluate(model, loader, loss_fn, device):
    model.eval()
    total_loss = 0.0
    total_examples = 0
    all_logits, all_targets = [], []

    with torch.inference_mode():
        for features, target in loader:
            features = features.to(device)
            target = target.to(device)
            logits = model(features)
            loss = loss_fn(logits, target)

            batch = target.shape[0]
            total_loss += loss.item() * batch
            total_examples += batch
            all_logits.append(logits.cpu())
            all_targets.append(target.cpu())

    return {
        "loss": total_loss / total_examples,
        "logits": torch.cat(all_logits),
        "targets": torch.cat(all_targets),
    }

This assumes that loss is computed as the mean per batch. If using reduction="sum", sample weights, token-level padding, or an incomplete final batch, the accumulation method must be adjusted accordingly. Metrics should be computed after aggregating predictions and should reflect the overall correct calculation, never simply average batch-level AUC/F1 values.

11. Train/Eval Mode and Validation

model.eval() does not mean disabling gradients; torch.inference_mode()/no_grad() also does not mean switching off dropout or BatchNorm. Validation typically involves both.

Validation is used for:

• early stopping and checkpoint selection;
• comparing training schedules and hyperparameters;
• calibration and threshold setting (requires appropriate data boundaries);
• identifying failure modes and analyzing learning curves.

Test data is not processed in every epoch loop. If test accuracy is checked every epoch, it has effectively become part of the validation process.

12. Checkpoint Must Be Able to Resume Training

Saving only weights is sufficient for inference, but not for seamless resumption of training. A proper training checkpoint typically includes:

• Model, optimizer, and scheduler states;
• AMP scaler state (if used);
• Epoch, global step, and best metric values;
• Sampler and data position states (required for precise recovery);
• Random number generator (RNG) states and seed policies;
• Configuration, code, dependencies, and data version information;
• Normalization, tokenizer, and label mapping configurations.

It's essential to save two checkpoints: "last" and "best validation," with clearly defined criteria for what constitutes "best," including the evaluation metric, direction of improvement, and tie-breaking rules. All checkpoint writes must be atomic and validated to ensure they can be successfully loaded and run inference, preventing the painful discovery of corrupted checkpoints after hours of training.

13. Reproducibility Has Boundaries

Setting Python/NumPy/PyTorch seeds only controls a subset of random sources. Hardware, parallel execution, nondeterministic kernels, and library versions can all alter results; bitwise consistency between CPU/GPU or versions is not guaranteed.

Deterministic algorithms may degrade performance or fail on operations without a deterministic implementation. Determinism can be beneficial during development and regression testing, but final benchmarking must document the observed behavior. Reporting multiple seeds or confidence ranges is generally more honest than showcasing just the best-performing seed.

14. Three Tests Before and After Training

Before Training

• shape/dtype/device;
• label range and mask;
• number of parameters, initial logits/loss;
• data duplication or leakage;
• finite forward/backward computation per single batch.

Overfitting on Small Datasets

Attempt to reduce the training loss to very low levels with just 1–2 batches. Failures typically indicate issues with implementation, model capacity, loss/label definition, or optimizer configuration; success does not guarantee generalization.

Formal Training

Record train/validation loss, task metrics, gradient and update norm, throughput, memory usage, learning rate, and checkpoint state. When anomalies are observed, first narrow down the reproducible case before applying any optimizer tricks.

Common Misconceptions

• All parameters cannot be initialized to 0: The key is breaking symmetry among weights at the same layer.
• Adam is the default optimal for any task: It's a strong candidate, but requires fair tuning and validation.
• BatchNorm must always be placed before activation: The order is part of architectural design.
• Larger batch sizes lead to faster convergence and worse generalization: Batch size must be balanced with update dynamics, training schedule, and computational cost.
• model.eval() has disabled autograd: Gradient recording and autograd mode are separate toggles.

Exercise

1. Observe the gradient symmetry caused by all-zero hidden weights.
2. Compare the activation and gradient variance across layers using Xavier versus Kaiming initialization.
3. Conduct a fair comparison under the same computational budget between SGD with Momentum and AdamW.
4. Intentionally forget eval(), and observe the validation fluctuations when Dropout or BatchNorm is present.
5. Save and restore the optimizer, scheduler, and RNG state to verify that the next update is reproducible.

Summary

Stable training arises from a set of interdependent choices: initialization preserves signal scale, the optimizer and learning-rate schedule determine updates, batch normalization alters statistics, and regularization constrains generalization. The training loop must explicitly manage mode, gradients, validation, and checkpoints.

The next chapter integrates these foundational elements into structured deep models: how CNNs encode locality, how RNNs maintain state, and how attention/Transformer architectures enable information exchange between positions based on contextual relevance.