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

Correct gradients do not guarantee stable training. A high learning rate can produce NaNs; changing the seed can change the result; validation may deteriorate while training loss keeps falling. Each symptom points to a different part of the training process.

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:

This fragment assumes model.hidden contains only Linear–ReLU hidden blocks. It initializes that submodule; the output head requires a choice suited to its own activation and loss.

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.hidden.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,{\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 gradient magnitudes differ across coordinates, 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. With the default track_running_stats=True, evaluation uses running statistics; disabling it keeps batch statistics in evaluation too. 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(1,\frac{c}{\Vert g\Vert }).$$

It limits the norm of the gradient per step and is commonly used in sequence models or unstable training scenarios. Global norm clipping multiplies a nonzero gradient by one positive scalar, preserving its direction. A threshold that is too small continually suppresses its magnitude. 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

This Python 3 / PyTorch fragment defines a full-precision, single-device loop. Supply a model already on device, nonempty loaders of (features, target) batches, an optimizer, and a scalar mean loss over equally weighted examples. It keeps all evaluation predictions in CPU memory.

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, error_if_nonfinite=True
        )
        optimizer.step()

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

    if total_examples == 0:
        raise ValueError("training loader must contain examples")
    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)
            if not torch.isfinite(loss):
                raise FloatingPointError(f"non-finite evaluation loss: {loss.item()}")

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

    if total_examples == 0:
        raise ValueError("evaluation loader must contain examples")
    return {
        "loss": total_loss / total_examples,
        "logits": torch.cat(all_logits),
        "targets": torch.cat(all_targets),
    }

Multiplying each batch mean by its size already handles a shorter final batch: means 2 and 8 over batches of sizes 3 and 1 give epoch loss 3.5, not 5. Class/sample weights, ignored targets, summed losses, or padded token losses may use a different denominator and require a matching accumulator. Aggregate predictions before computing non-additive metrics such as AUC or F1.

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.